Spaces:
Running on Zero
Running on Zero
File size: 4,412 Bytes
7405e85 | 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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | 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
@torch.no_grad()
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()
@cache
def get_lpips(device: torch.device) -> LPIPS:
model = LPIPS(net="vgg").eval().requires_grad_(False)
return model.to(device=device, dtype=torch.float32)
@torch.no_grad()
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()
@torch.no_grad()
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
|