File size: 2,944 Bytes
251713e | 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 | """Evaluation metrics for MAVT.
Image : rFID (requires torchmetrics-image), PSNR, SSIM
Video : temporal PSNR (per-frame then averaged)
3D : per-plane PSNR, SSIM
"""
from __future__ import annotations
from typing import Dict, Optional
import torch
import torch.nn.functional as F
def psnr(pred: torch.Tensor, target: torch.Tensor, max_val: float = 1.0) -> torch.Tensor:
"""PSNR in dB. Inputs expected in [0, 1]."""
mse = F.mse_loss(pred, target)
return 10 * torch.log10(max_val ** 2 / (mse + 1e-8))
def ssim(pred: torch.Tensor, target: torch.Tensor, window_size: int = 11) -> torch.Tensor:
"""Simplified SSIM via torchmetrics if available, else MSE proxy."""
try:
from torchmetrics.functional.image import structural_similarity_index_measure as _ssim
return _ssim(pred, target, data_range=1.0)
except ImportError:
return 1.0 - F.mse_loss(pred, target)
def compute_image_metrics(pred: torch.Tensor, target: torch.Tensor) -> Dict[str, float]:
"""pred, target: (B, 3, H, W) in [0, 1]."""
return {
'psnr': psnr(pred, target).item(),
'ssim': ssim(pred, target).item(),
}
def compute_video_metrics(pred: torch.Tensor, target: torch.Tensor) -> Dict[str, float]:
"""pred, target: (B, 3, T, H, W) in [0, 1]."""
T = pred.shape[2]
psnr_vals = [psnr(pred[:, :, t], target[:, :, t]).item() for t in range(T)]
return {
'temporal_psnr': sum(psnr_vals) / T,
'temporal_psnr_min': min(psnr_vals),
}
def compute_threed_metrics(pred: torch.Tensor, target: torch.Tensor) -> Dict[str, float]:
"""pred, target: (B, 3, 3, H, W) — 3 planes each (B, 3, H, W)."""
plane_names = ['xy', 'xz', 'yz']
metrics: Dict[str, float] = {}
for i, name in enumerate(plane_names):
metrics[f'psnr_{name}'] = psnr(pred[:, i], target[:, i]).item()
metrics[f'ssim_{name}'] = ssim(pred[:, i], target[:, i]).item()
metrics['psnr_mean'] = sum(metrics[f'psnr_{n}'] for n in plane_names) / 3
return metrics
class FIDTracker:
"""Accumulates features for FID computation using torchmetrics."""
def __init__(self, feature: int = 2048):
try:
from torchmetrics.image.fid import FrechetInceptionDistance
self._fid = FrechetInceptionDistance(feature=feature, normalize=True)
self._available = True
except ImportError:
self._available = False
def update_real(self, imgs: torch.Tensor) -> None:
if self._available:
self._fid.update(imgs.clamp(0, 1), real=True)
def update_fake(self, imgs: torch.Tensor) -> None:
if self._available:
self._fid.update(imgs.clamp(0, 1), real=False)
def compute(self) -> Optional[float]:
if self._available:
return float(self._fid.compute())
return None
def reset(self) -> None:
if self._available:
self._fid.reset()
|