metrics_eval.py

import math
from typing import Dict, Optional, Tuple, Any

import torch


def _to_01(x: torch.Tensor) -> torch.Tensor:
    if x.min() < 0:
        x = (x + 1.0) / 2.0
    return x.clamp(0.0, 1.0)


def _ensure_3ch(x: torch.Tensor) -> torch.Tensor:
    if x.size(1) == 1:
        return x.repeat(1, 3, 1, 1)
    return x


def _psnr(recons: torch.Tensor, targets: torch.Tensor) -> float:
    mse = torch.mean((recons - targets) ** 2, dim=(1, 2, 3))
    mse = torch.clamp(mse, min=1e-10)
    psnr = 20.0 * torch.log10(torch.tensor(1.0, device=recons.device)) - 10.0 * torch.log10(mse)
    return psnr.mean().item()


class MetricsEvaluator:
    def __init__(self, device: torch.device):
        self.device = device

        self._has_torchmetrics = False
        self._has_piq = False
        self._has_lpips = False

        try:
            from torchmetrics.image.ssim import StructuralSimilarityIndexMeasure
            from torchmetrics.image.ms_ssim import MultiScaleStructuralSimilarityIndexMeasure
            from torchmetrics.image.fid import FrechetInceptionDistance
            from torchmetrics.image.kid import KernelInceptionDistance
            from torchmetrics.image.inception import InceptionScore
            self.StructuralSimilarityIndexMeasure = StructuralSimilarityIndexMeasure
            self.MultiScaleStructuralSimilarityIndexMeasure = MultiScaleStructuralSimilarityIndexMeasure
            self.FrechetInceptionDistance = FrechetInceptionDistance
            self.KernelInceptionDistance = KernelInceptionDistance
            self.InceptionScore = InceptionScore
            self._has_torchmetrics = True
        except Exception:
            self._has_torchmetrics = False

        try:
            import piq
            self.piq = piq
            self._has_piq = True
        except Exception:
            self._has_piq = False

        try:
            import lpips
            self.lpips = lpips
            self._has_lpips = True
        except Exception:
            self._has_lpips = False

    def compute(self,
                recons: torch.Tensor,
                targets: torch.Tensor,
                generated: Optional[torch.Tensor] = None) -> Tuple[Dict[str, Any], Dict[str, str]]:
        results: Dict[str, Any] = {}
        errors: Dict[str, str] = {}

        recons = _to_01(recons.to(self.device))
        targets = _to_01(targets.to(self.device))
        recons_3 = _ensure_3ch(recons)
        targets_3 = _ensure_3ch(targets)

        if generated is None:
            generated = recons_3
        else:
            generated = _ensure_3ch(_to_01(generated.to(self.device)))

        # PSNR
        try:
            results["PSNR"] = _psnr(recons_3, targets_3)
        except Exception as exc:
            errors["PSNR"] = str(exc)

        # SSIM / MS-SSIM
        if self._has_torchmetrics:
            try:
                ssim = self.StructuralSimilarityIndexMeasure(data_range=1.0).to(self.device)
                results["SSIM"] = float(ssim(recons_3, targets_3).item())
            except Exception as exc:
                errors["SSIM"] = str(exc)
            try:
                ms_ssim = self.MultiScaleStructuralSimilarityIndexMeasure(data_range=1.0).to(self.device)
                results["MS-SSIM"] = float(ms_ssim(recons_3, targets_3).item())
            except Exception as exc:
                errors["MS-SSIM"] = str(exc)
        elif self._has_piq:
            try:
                results["SSIM"] = float(self.piq.ssim(recons_3, targets_3, data_range=1.0).item())
            except Exception as exc:
                errors["SSIM"] = str(exc)
            try:
                results["MS-SSIM"] = float(self.piq.multi_scale_ssim(recons_3, targets_3, data_range=1.0).item())
            except Exception as exc:
                errors["MS-SSIM"] = str(exc)
        else:
            errors["SSIM"] = "torchmetrics or piq not installed"
            errors["MS-SSIM"] = "torchmetrics or piq not installed"

        # LPIPS
        if self._has_lpips:
            try:
                lpips_model = self.lpips.LPIPS(net="alex").to(self.device)
                recons_lp = recons_3 * 2.0 - 1.0
                targets_lp = targets_3 * 2.0 - 1.0
                results["LPIPS"] = float(lpips_model(recons_lp, targets_lp).mean().item())
            except Exception as exc:
                errors["LPIPS"] = str(exc)
        else:
            errors["LPIPS"] = "lpips not installed"

        # VIF
        if self._has_piq:
            try:
                results["VIF"] = float(self.piq.vif_p(recons_3, targets_3, data_range=1.0).item())
            except Exception as exc:
                errors["VIF"] = str(exc)
        else:
            errors["VIF"] = "piq not installed"

        # FID / KID / IS
        if self._has_torchmetrics:
            try:
                fid = self.FrechetInceptionDistance(feature=2048, normalize=True).to(self.device)
                fid.update(targets_3, real=True)
                fid.update(generated, real=False)
                results["FID"] = float(fid.compute().item())
            except Exception as exc:
                errors["FID"] = str(exc)
            try:
                kid = self.KernelInceptionDistance(subset_size=50, normalize=True).to(self.device)
                kid.update(targets_3, real=True)
                kid.update(generated, real=False)
                kid_mean, kid_std = kid.compute()
                results["KID_mean"] = float(kid_mean.item())
                results["KID_std"] = float(kid_std.item())
            except Exception as exc:
                errors["KID"] = str(exc)
            try:
                is_metric = self.InceptionScore(normalize=True).to(self.device)
                is_metric.update(generated)
                is_mean, is_std = is_metric.compute()
                results["IS_mean"] = float(is_mean.item())
                results["IS_std"] = float(is_std.item())
            except Exception as exc:
                errors["IS"] = str(exc)
        else:
            errors["FID"] = "torchmetrics not installed"
            errors["KID"] = "torchmetrics not installed"
            errors["IS"] = "torchmetrics not installed"

        return results, errors