58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
"""Windowed SSIM with no scipy/skimage dependency.
|
||
|
||
Implements the Wang et al. structural-similarity index using a uniform window
|
||
(box filter) computed via an integral image — fast, vectorised, and dependency
|
||
-light (numpy only). Good enough to catch rendering regressions while tolerating
|
||
sub-pixel antialiasing noise.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import numpy as np
|
||
|
||
_C1 = (0.01 * 255) ** 2
|
||
_C2 = (0.03 * 255) ** 2
|
||
|
||
|
||
def _mean_filter(img: np.ndarray, k: int) -> np.ndarray:
|
||
"""Box-mean over a k×k window, same shape, edge-clamped window area."""
|
||
h, w = img.shape
|
||
pad = k // 2
|
||
integral = np.zeros((h + 1, w + 1), dtype=np.float64)
|
||
integral[1:, 1:] = img.cumsum(0).cumsum(1)
|
||
|
||
ys = np.clip(np.arange(h) - pad, 0, h)
|
||
ye = np.clip(np.arange(h) + pad + 1, 0, h)
|
||
xs = np.clip(np.arange(w) - pad, 0, w)
|
||
xe = np.clip(np.arange(w) + pad + 1, 0, w)
|
||
|
||
a = integral[ye][:, xe]
|
||
b = integral[ye][:, xs]
|
||
c = integral[ys][:, xe]
|
||
d = integral[ys][:, xs]
|
||
total = a - b - c + d
|
||
area = (ye - ys)[:, None] * (xe - xs)[None, :]
|
||
return total / area
|
||
|
||
|
||
def ssim(a: np.ndarray, b: np.ndarray, k: int = 7) -> float:
|
||
"""Mean SSIM in [0, 1]. Mismatched shapes score 0.0 (always a regression)."""
|
||
if a.shape != b.shape:
|
||
return 0.0
|
||
a = a.astype(np.float64)
|
||
b = b.astype(np.float64)
|
||
|
||
mu_a = _mean_filter(a, k)
|
||
mu_b = _mean_filter(b, k)
|
||
mu_a2 = mu_a * mu_a
|
||
mu_b2 = mu_b * mu_b
|
||
mu_ab = mu_a * mu_b
|
||
|
||
var_a = _mean_filter(a * a, k) - mu_a2
|
||
var_b = _mean_filter(b * b, k) - mu_b2
|
||
cov_ab = _mean_filter(a * b, k) - mu_ab
|
||
|
||
num = (2 * mu_ab + _C1) * (2 * cov_ab + _C2)
|
||
den = (mu_a2 + mu_b2 + _C1) * (var_a + var_b + _C2)
|
||
return float(np.clip(num / den, 0.0, 1.0).mean())
|