File size: 7,890 Bytes
223873d bed503a 223873d bed503a 223873d bed503a 223873d bed503a 223873d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 | """Shared helpers for codec test / eval scripts."""
from __future__ import annotations
import os
from typing import Dict, Optional, Tuple
import torch
from flexicm.utils.alignment import Alignment
from flexicm.utils.train_utils import AverageMeter
def resolve_ckpt(path: str, repo_root: str, label: str = "checkpoint") -> str:
if not path:
raise FileNotFoundError(f"{label}: path is empty")
if not os.path.isabs(path):
path = os.path.join(repo_root, path)
placeholder = os.path.join(os.path.dirname(path), "PLACEHOLDER")
if os.path.isfile(placeholder) and not os.path.isfile(path):
raise FileNotFoundError(
f"{label}: not ready (PLACEHOLDER still present):\n {placeholder}\n"
f"Expected real weights at:\n {path}\n"
"See checkpoints/README.md"
)
if not os.path.isfile(path):
raise FileNotFoundError(f"{label}: missing file: {path}")
if os.path.basename(path) == "PLACEHOLDER" or path.endswith(".txt"):
raise FileNotFoundError(f"{label}: refusing placeholder file: {path}")
return path
def crop_feature_to_image(h: torch.Tensor, image_hw: Tuple[int, int]) -> torch.Tensor:
"""Crop decoded feature h (H/4, W/4 of padded input) to original image size / 4."""
H, W = image_hw
return h[..., : H // 4, : W // 4]
def pad_for_codec(images: torch.Tensor, divisor: int = 256, device=None):
align = Alignment(divisor=divisor, mode="pad", padding_mode="constant")
if device is not None:
align = align.to(device)
return align.align(images), align
def likelihood_bpp(likelihoods: Dict[str, torch.Tensor], num_pixels: int) -> torch.Tensor:
import math
return sum(
(torch.log(lik).sum() / (-math.log(2) * num_pixels))
for lik in likelihoods.values()
)
def actual_bitstream_bpp(strings, num_pixels: int) -> float:
"""Estimate bpp from CompressAI byte strings: [[y_bytes...], [z_bytes...]]."""
total_bits = 0
for group in strings:
for s in group:
if isinstance(s, (bytes, bytearray)):
total_bits += len(s) * 8
elif torch.is_tensor(s):
total_bits += int(s.numel() * s.element_size() * 8)
else:
total_bits += len(s) * 8
return total_bits / float(num_pixels)
@torch.no_grad()
def test_taic_loader(
model,
teacher,
loader,
criterion,
device,
align_divisor: int = 256,
run_actual_bpp: bool = False,
max_batches: Optional[int] = None,
log_every: int = 50,
):
"""Run codec test: likelihood bpp + feature distortion (+ optional real bpp)."""
model.eval()
teacher.eval()
meters = {k: AverageMeter() for k in ("loss", "bpp", "distortion", "actual_bpp")}
if run_actual_bpp:
model.update(force=True)
for i, images in enumerate(loader):
if max_batches is not None and i >= max_batches:
break
images = images.to(device)
N, _, H, W = images.shape
num_pixels = N * H * W
# Keep codec + teacher on the padded grid so Swin always sees a
# patch/window-divisible size. bpp still uses the original pixel count.
x, _ = pad_for_codec(images, divisor=align_divisor, device=device)
out = model(x)
gt = teacher.gt_features(x)
pred = teacher.pred_features(out["h"])
stats = criterion(out, pred, gt, num_pixels=num_pixels)
meters["loss"].update(stats["loss"].item(), n=N)
meters["bpp"].update(stats["bpp"].item(), n=N)
meters["distortion"].update(stats["distortion"].item(), n=N)
if run_actual_bpp:
try:
enc = model.compress(x)
dec = model.decompress(
enc["strings"], enc["shape"], x_size=(x.shape[2], x.shape[3])
)
abpp = actual_bitstream_bpp(enc["strings"], num_pixels)
meters["actual_bpp"].update(abpp, n=N)
# sanity: decoded h spatial size
_ = dec["h"]
except Exception as e:
if i == 0:
print(f"[warn] actual bpp / compress-decompress failed: {e}")
if i % log_every == 0:
msg = (
f"[{i}/{len(loader)}] bpp={meters['bpp'].avg:.4f} "
f"D={meters['distortion'].avg:.6f} loss={meters['loss'].avg:.4f}"
)
if run_actual_bpp and meters["actual_bpp"].count > 0:
msg += f" actual_bpp={meters['actual_bpp'].avg:.4f}"
print(msg)
result = {
"bpp": meters["bpp"].avg,
"distortion": meters["distortion"].avg,
"loss": meters["loss"].avg,
"num_batches": meters["bpp"].count,
}
if run_actual_bpp and meters["actual_bpp"].count > 0:
result["actual_bpp"] = meters["actual_bpp"].avg
return result
@torch.no_grad()
def test_ctaic_loader(
ext_model,
base_model,
teacher,
loader,
criterion,
device,
use_condition: bool = True,
align_divisor: int = 256,
run_actual_bpp: bool = False,
max_batches: Optional[int] = None,
log_every: int = 50,
):
"""Codec test for C-TAIC; bpp is extension-layer only (paper Sec.IV.E.2)."""
ext_model.eval()
base_model.eval()
teacher.eval()
meters = {k: AverageMeter() for k in ("loss", "bpp", "distortion", "actual_bpp")}
if run_actual_bpp:
ext_model.update(force=True)
for i, images in enumerate(loader):
if max_batches is not None and i >= max_batches:
break
images = images.to(device)
N, _, H, W = images.shape
num_pixels = N * H * W
# Keep codec + teacher on the padded grid (same rationale as TAIC test).
x, _ = pad_for_codec(images, divisor=align_divisor, device=device)
y_b = None
if use_condition:
base_out = base_model(x)
y_b = base_out["y_hat"]
out = ext_model(x, y_b_hat=y_b, use_condition=use_condition and y_b is not None)
gt = teacher.gt_features(x)
pred = teacher.pred_features(out["h"])
stats = criterion(out, pred, gt, num_pixels=num_pixels)
meters["loss"].update(stats["loss"].item(), n=N)
meters["bpp"].update(stats["bpp"].item(), n=N)
meters["distortion"].update(stats["distortion"].item(), n=N)
if run_actual_bpp:
try:
enc = ext_model.compress(
x, y_b_hat=y_b, use_condition=use_condition and y_b is not None
)
dec = ext_model.decompress(
enc["strings"],
enc["shape"],
x_size=(x.shape[2], x.shape[3]),
y_b_hat=y_b,
use_condition=use_condition and y_b is not None,
)
abpp = actual_bitstream_bpp(enc["strings"], num_pixels)
meters["actual_bpp"].update(abpp, n=N)
_ = dec["h"]
except Exception as e:
if i == 0:
print(f"[warn] actual bpp / compress-decompress failed: {e}")
if i % log_every == 0:
msg = (
f"[{i}/{len(loader)}] bpp={meters['bpp'].avg:.4f} "
f"D={meters['distortion'].avg:.6f} loss={meters['loss'].avg:.4f}"
)
if run_actual_bpp and meters["actual_bpp"].count > 0:
msg += f" actual_bpp={meters['actual_bpp'].avg:.4f}"
print(msg)
result = {
"bpp": meters["bpp"].avg,
"distortion": meters["distortion"].avg,
"loss": meters["loss"].avg,
"num_batches": meters["bpp"].count,
"use_condition": use_condition,
}
if run_actual_bpp and meters["actual_bpp"].count > 0:
result["actual_bpp"] = meters["actual_bpp"].avg
return result
|