Spaces:
Running on Zero
Running on Zero
| from functools import cache | |
| import torch | |
| import torch.nn.functional as F | |
| from einops import reduce | |
| from jaxtyping import Float | |
| from lpips import LPIPS | |
| from torch import Tensor | |
| def compute_psnr( | |
| ground_truth: Float[Tensor, "batch channel height width"], | |
| predicted: Float[Tensor, "batch channel height width"], | |
| ) -> Float[Tensor, " batch"]: | |
| with torch.autocast(device_type=predicted.device.type, enabled=False): | |
| ground_truth = ground_truth.float().clip(min=0, max=1) | |
| predicted = predicted.float().clip(min=0, max=1) | |
| mse = reduce((ground_truth - predicted) ** 2, "b c h w -> b", "mean") | |
| return -10 * mse.log10() | |
| def get_lpips(device: torch.device) -> LPIPS: | |
| model = LPIPS(net="vgg").eval().requires_grad_(False) | |
| return model.to(device=device, dtype=torch.float32) | |
| def compute_lpips( | |
| ground_truth: Float[Tensor, "batch channel height width"], | |
| predicted: Float[Tensor, "batch channel height width"], | |
| ) -> Float[Tensor, " batch"]: | |
| with torch.autocast(device_type=predicted.device.type, enabled=False): | |
| value = get_lpips(predicted.device).forward( | |
| ground_truth.float(), | |
| predicted.float(), | |
| normalize=True, | |
| ) | |
| return value[:, 0, 0, 0].float() | |
| def compute_ssim( | |
| ground_truth: Float[Tensor, "batch channel height width"], | |
| predicted: Float[Tensor, "batch channel height width"], | |
| *, | |
| win_size: int = 11, | |
| gaussian_weights: bool = True, | |
| data_range: float = 1.0, | |
| use_sample_covariance: bool = True, | |
| sigma: float = 1.5, | |
| k1: float = 0.01, | |
| k2: float = 0.03, | |
| ) -> Float[Tensor, " batch"]: | |
| if ground_truth.shape != predicted.shape or ground_truth.ndim != 4: | |
| raise ValueError("SSIM inputs must have the same BCHW shape.") | |
| if win_size <= 0 or win_size % 2 == 0: | |
| raise ValueError("win_size must be a positive odd integer.") | |
| if min(ground_truth.shape[-2:]) < win_size: | |
| raise ValueError("win_size must not exceed the image dimensions.") | |
| if gaussian_weights and sigma <= 0: | |
| raise ValueError("sigma must be positive when using Gaussian weights.") | |
| if data_range <= 0 or k1 < 0 or k2 < 0: | |
| raise ValueError("data_range must be positive and k1/k2 non-negative.") | |
| with torch.autocast(device_type=predicted.device.type, enabled=False): | |
| ground_truth = ground_truth.float() | |
| predicted = predicted.float() | |
| if gaussian_weights: | |
| coordinates = torch.arange( | |
| win_size, | |
| dtype=torch.float32, | |
| device=predicted.device, | |
| ) | |
| coordinates -= (win_size - 1) / 2 | |
| kernel_1d = torch.exp(-(coordinates.square()) / (2 * sigma**2)) | |
| kernel_1d /= kernel_1d.sum() | |
| kernel = torch.outer(kernel_1d, kernel_1d) | |
| else: | |
| kernel = torch.full( | |
| (win_size, win_size), | |
| 1 / win_size**2, | |
| dtype=torch.float32, | |
| device=predicted.device, | |
| ) | |
| channels = predicted.shape[1] | |
| moments = torch.cat( | |
| ( | |
| ground_truth, | |
| predicted, | |
| ground_truth.square(), | |
| predicted.square(), | |
| ground_truth * predicted, | |
| ), | |
| dim=1, | |
| ) | |
| kernels = kernel[None, None].expand(5 * channels, 1, -1, -1) | |
| ux, uy, uxx, uyy, uxy = F.conv2d( | |
| moments, | |
| kernels, | |
| groups=5 * channels, | |
| ).chunk(5, dim=1) | |
| covariance_normalization = 1.0 | |
| if use_sample_covariance: | |
| num_pixels = win_size**2 | |
| if num_pixels == 1: | |
| raise ValueError( | |
| "Sample covariance requires a window larger than one pixel." | |
| ) | |
| covariance_normalization = num_pixels / (num_pixels - 1) | |
| vx = covariance_normalization * (uxx - ux.square()) | |
| vy = covariance_normalization * (uyy - uy.square()) | |
| vxy = covariance_normalization * (uxy - ux * uy) | |
| c1 = (k1 * data_range) ** 2 | |
| c2 = (k2 * data_range) ** 2 | |
| ssim = ( | |
| (2 * ux * uy + c1) | |
| * (2 * vxy + c2) | |
| / ((ux.square() + uy.square() + c1) * (vx + vy + c2)) | |
| ) | |
| ssim = ssim.mean(dim=(1, 2, 3)) | |
| return ssim | |