| """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)) |
| |
| s = s[..., pad:-pad, pad:-pad] |
| return s.mean(dim=(1, 2, 3)) |
|
|
| @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)} |
|
|
