TorchDiff 2.0.0__py3-none-any.whl

unclip/val_metrics.py

@@ -0,0 +1,221 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.attention import sdpa_kernel, SDPBackend
+from pytorch_fid import fid_score
+from torchvision.utils import save_image
+from transformers import BertModel
+import os
+import lpips
+import math
+import shutil
+from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+from torchmetrics.image.fid import FrechetInceptionDistance
+from torchvision.utils import save_image
+from typing import Optional, Tuple, List
+
+
+
+
+
+
+
+class Metrics:
+    """Computes image quality metrics for evaluating diffusion models.
+
+    Supports Mean Squared Error (MSE), Peak Signal-to-Noise Ratio (PSNR), Structural
+    Similarity Index (SSIM), Fréchet Inception Distance (FID), and Learned Perceptual
+    Image Patch Similarity (LPIPS) for comparing generated and ground truth images.
+
+    Parameters
+    ----------
+    device : str, optional
+        Device for computation (e.g., 'cuda', 'cpu') (default: 'cuda').
+    fid : bool, optional
+        If True, compute FID score (default: True).
+    metrics : bool, optional
+        If True, compute MSE, PSNR, and SSIM (default: False).
+    lpips : bool, optional
+        If True, compute LPIPS using VGG backbone (default: False).
+    """
+
+    def __init__(
+            self,
+            device: str = "cuda",
+            fid: bool = True,
+            metrics: bool = False,
+            lpips_: bool = False
+    ) -> None:
+        self.device = device
+        self.fid = fid
+        self.metrics = metrics
+        self.lpips = lpips_
+        self.lpips_model = LearnedPerceptualImagePatchSimilarity(
+            net_type='vgg',
+            normalize=True  # This handles [0,1] -> [-1,1] conversion
+        ).to(device) if self.lpips else None
+        self.temp_dir_real = "temp_real"
+        self.temp_dir_fake = "temp_fake"
+
+    def compute_fid(self, real_images: torch.Tensor, fake_images: torch.Tensor) -> float:
+        """Computes the Fréchet Inception Distance (FID) between real and generated images.
+
+        Saves images to temporary directories and uses Inception V3 to compute FID,
+        cleaning up directories afterward.
+
+        Parameters
+        ----------
+        real_images : torch.Tensor
+            Real images, shape (batch_size, channels, height, width), in [-1, 1].
+        fake_images : torch.Tensor
+            Generated images, same shape, in [-1, 1].
+
+        Returns
+        -------
+        fid (float) - FID score, or `float('inf')` if computation fails.
+
+        **Notes**
+
+        - Images are normalized to [0, 1] and saved as PNG files for FID computation.
+        - Uses Inception V3 with 2048-dimensional features (`dims=2048`).
+        """
+        if real_images.shape != fake_images.shape:
+            raise ValueError(f"Shape mismatch: real_images {real_images.shape}, fake_images {fake_images.shape}")
+
+        real_images = (real_images + 1) / 2
+        fake_images = (fake_images + 1) / 2
+        real_images = real_images.clamp(0, 1).cpu()
+        fake_images = fake_images.clamp(0, 1).cpu()
+
+        os.makedirs(self.temp_dir_real, exist_ok=True)
+        os.makedirs(self.temp_dir_fake, exist_ok=True)
+
+        try:
+            for i, (real, fake) in enumerate(zip(real_images, fake_images)):
+                save_image(real, f"{self.temp_dir_real}/{i}.png")
+                save_image(fake, f"{self.temp_dir_fake}/{i}.png")
+
+            fid = fid_score.calculate_fid_given_paths(
+                paths=[self.temp_dir_real, self.temp_dir_fake],
+                batch_size=50,
+                device=self.device,
+                dims=2048
+            )
+        except Exception as e:
+            print(f"Error computing FID: {e}")
+            fid = float('inf')
+        finally:
+            shutil.rmtree(self.temp_dir_real, ignore_errors=True)
+            shutil.rmtree(self.temp_dir_fake, ignore_errors=True)
+
+        return fid
+
+    def compute_metrics(self, x: torch.Tensor, x_hat: torch.Tensor) -> Tuple[float, float, float]:
+        """Computes MSE, PSNR, and SSIM for evaluating image quality.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Ground truth images, shape (batch_size, channels, height, width).
+        x_hat : torch.Tensor
+            Generated images, same shape as `x`.
+
+        Returns
+        -------
+        mse : float
+            Mean squared error.
+        psnr : float
+            Peak signal-to-noise ratio.
+        ssim : float
+            Structural similarity index (mean over batch).
+        """
+        if x.shape != x_hat.shape:
+            raise ValueError(f"Shape mismatch: x {x.shape}, x_hat {x_hat.shape}")
+
+        mse = F.mse_loss(x_hat, x)
+        psnr = -10 * torch.log10(mse)
+        c1, c2 = (0.01 * 2) ** 2, (0.03 * 2) ** 2  # Adjusted for [-1, 1] range
+        eps = 1e-8
+        mu_x = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1)
+        mu_y = F.avg_pool2d(x_hat, kernel_size=3, stride=1, padding=1)
+        mu_xy = mu_x * mu_y
+        sigma_x_sq = F.avg_pool2d(x.pow(2), kernel_size=3, stride=1, padding=1) - mu_x.pow(2)
+        sigma_y_sq = F.avg_pool2d(x_hat.pow(2), kernel_size=3, stride=1, padding=1) - mu_y.pow(2)
+        sigma_xy = F.avg_pool2d(x * x_hat, kernel_size=3, stride=1, padding=1) - mu_xy
+        ssim = ((2 * mu_xy + c1) * (2 * sigma_xy + c2)) / (
+            (mu_x.pow(2) + mu_y.pow(2) + c1) * (sigma_x_sq + sigma_y_sq + c2) + eps
+        )
+
+        return mse.item(), psnr.item(), ssim.mean().item()
+
+    def compute_lpips(self, x: torch.Tensor, x_hat: torch.Tensor) -> float:
+        """Computes LPIPS using a pre-trained VGG network.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Ground truth images, shape (batch_size, channels, height, width), in [-1, 1].
+        x_hat : torch.Tensor
+            Generated images, same shape as `x`.
+
+        Returns
+        -------
+        lpips (float) - Mean LPIPS score over the batch.
+        """
+        if self.lpips_model is None:
+            raise RuntimeError("LPIPS model not initialized; set lpips=True in __init__")
+        if x.shape != x_hat.shape:
+            raise ValueError(f"Shape mismatch: x {x.shape}, x_hat {x_hat.shape}")
+
+        # Normalize inputs to [0, 1] range
+        x = (x + 1) / 2  # Convert from [-1, 1] to [0, 1]
+        x_hat = (x_hat + 1) / 2
+        x = x.clamp(0, 1)  # Ensure values are in [0, 1]
+        x_hat = x_hat.clamp(0, 1)
+
+        x = x.to(self.device)
+        x_hat = x_hat.to(self.device)
+
+        # Convert grayscale to RGB if needed
+        if x.shape[1] == 1:
+            x = x.repeat(1, 3, 1, 1)  # Repeat grayscale channel 3 times
+        if x_hat.shape[1] == 1:
+            x_hat = x_hat.repeat(1, 3, 1, 1)
+
+        return self.lpips_model(x, x_hat).mean().item()
+
+    def forward(self, x: torch.Tensor, x_hat: torch.Tensor) -> Tuple[float, float, float, float, float]:
+        """Computes specified metrics for ground truth and generated images.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Ground truth images, shape (batch_size, channels, height, width), in [-1, 1].
+        x_hat : torch.Tensor
+            Generated images, same shape as `x`.
+
+        Returns
+        -------
+        fid : float, or `float('inf')` if not computed
+            Mean FID score.
+        mse : float, or None if not computed
+            Mean MSE
+        psnr : float, or None if not computed
+             Mean PSNR
+        ssim : float, or None if not computed
+            Mean SSIM
+        lpips_score :  float, or None if not computed
+            Mean LPIPS score
+        """
+        fid = float('inf')
+        mse, psnr, ssim = None, None, None
+        lpips_score = None
+
+        if self.metrics:
+            mse, psnr, ssim = self.compute_metrics(x, x_hat)
+        if self.fid:
+            fid = self.compute_fid(x, x_hat)
+        if self.lpips:
+            lpips_score = self.compute_lpips(x, x_hat)
+
+        return fid, mse, psnr, ssim, lpips_score