File size: 3,645 Bytes
f70ac4f | 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 | """PSNR / SSIM / LPIPS between a strategy's frames and its matched FFFF frames.
Computed on the decoded RGB tensors in [0, 1], before MP4 encoding, as the
protocol requires (section 8.1) -- video compression noise would otherwise leak
into the numbers.
SSIM is implemented here rather than taken from skimage: skimage needs ~30 ms per
480x832 frame on CPU, and this evaluation compares ~500k frames.
"""
import torch
import torch.nn.functional as F
def _gaussian_window(size=11, sigma=1.5, device=None, dtype=torch.float32):
coords = torch.arange(size, device=device, dtype=dtype) - (size - 1) / 2.0
g = torch.exp(-(coords ** 2) / (2 * sigma ** 2))
g = g / g.sum()
return g[:, None] @ g[None, :]
class PixelMetrics:
"""Frame-wise metrics against a reference clip. Call once per (strategy, prompt)."""
def __init__(self, device, lpips_net="alex", ssim_window=11, ssim_sigma=1.5):
self.device = device
self.window = _gaussian_window(ssim_window, ssim_sigma, device=device)[None, None]
self.window_size = ssim_window
import lpips
self.lpips = lpips.LPIPS(net=lpips_net).to(device).eval()
for p in self.lpips.parameters():
p.requires_grad_(False)
def ssim(self, x, y):
"""x, y: [T, C, H, W] in [0, 1]. Returns the mean over frames and channels."""
t, c, h, w = x.shape
win = self.window.expand(c, 1, self.window_size, self.window_size)
pad = self.window_size // 2
mu_x = F.conv2d(x, win, padding=pad, groups=c)
mu_y = F.conv2d(y, win, padding=pad, groups=c)
mu_x2, mu_y2, mu_xy = mu_x * mu_x, mu_y * mu_y, mu_x * mu_y
sig_x = F.conv2d(x * x, win, padding=pad, groups=c) - mu_x2
sig_y = F.conv2d(y * y, win, padding=pad, groups=c) - mu_y2
sig_xy = F.conv2d(x * y, win, padding=pad, groups=c) - mu_xy
c1, c2 = 0.01 ** 2, 0.03 ** 2
s = ((2 * mu_xy + c1) * (2 * sig_xy + c2)) / ((mu_x2 + mu_y2 + c1) * (sig_x + sig_y + c2))
# Drop the padded border, as skimage does with its default crop.
s = s[..., pad:-pad, pad:-pad]
return s.mean(dim=(1, 2, 3)) # per frame
@torch.no_grad()
def compute(self, frames, ref, chunk=16):
"""frames, ref: [T, C, H, W] float in [0, 1] on any device.
Returns per-prompt mean MSE (for the protocol's ratio-of-means PSNR),
plus mean SSIM and mean LPIPS over frames.
"""
frames = frames.to(self.device, torch.float32)
ref = ref.to(self.device, torch.float32)
assert frames.shape == ref.shape, (frames.shape, ref.shape)
mse_terms, ssim_terms, lpips_terms = [], [], []
for i in range(0, frames.shape[0], chunk):
a, b = frames[i:i + chunk], ref[i:i + chunk]
mse_terms.append(((a - b) ** 2).mean(dim=(1, 2, 3)))
ssim_terms.append(self.ssim(a, b))
lpips_terms.append(self.lpips(a * 2 - 1, b * 2 - 1).flatten())
mse = torch.cat(mse_terms)
ssim = torch.cat(ssim_terms)
lp = torch.cat(lpips_terms)
mean_mse = mse.mean().item()
return {
"mean_mse": mean_mse,
"psnr": -10.0 * torch.log10(torch.tensor(max(mean_mse, 1e-12))).item(),
"ssim": ssim.mean().item(),
"lpips": lp.mean().item(),
"num_frames": int(frames.shape[0]),
}
@staticmethod
def identity(num_frames):
"""Fixed self-comparison values the protocol pins for FFFF vs itself."""
return {"mean_mse": 1e-12, "psnr": 120.0, "ssim": 1.0, "lpips": 0.0,
"num_frames": int(num_frames)}
|