File size: 3,160 Bytes
8bfc737 | 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 | from __future__ import annotations
import torch
CSI_THRESHOLDS = torch.tensor([16, 74, 133, 160, 181, 219], dtype=torch.float32) / 255.0
def mse(prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
return torch.mean((prediction - target) ** 2)
def _ssim_per_sample(prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
prediction, target = prediction.clamp(0.0, 1.0), target.clamp(0.0, 1.0)
mu_x, mu_y = prediction.mean((2, 3)), target.mean((2, 3))
centered_x = prediction - mu_x[:, :, None, None, :]
centered_y = target - mu_y[:, :, None, None, :]
var_x = centered_x.square().mean((2, 3))
var_y = centered_y.square().mean((2, 3))
covariance = (centered_x * centered_y).mean((2, 3))
c1, c2 = 0.01**2, 0.03**2
score = ((2 * mu_x * mu_y + c1) * (2 * covariance + c2)) / (
(mu_x.square() + mu_y.square() + c1) * (var_x + var_y + c2)
)
return score.mean((1, 2))
def metric_sums(prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
"""Return additive sample-level MSE/MAE/SSIM sums and CSI event counts."""
prediction, target = prediction.float(), target.float()
reduce_dims = tuple(range(1, prediction.ndim))
values = [
(prediction - target).square().mean(reduce_dims).sum(),
(prediction - target).abs().mean(reduce_dims).sum(),
_ssim_per_sample(prediction, target).sum(),
prediction.new_tensor(prediction.shape[0]),
]
for threshold in CSI_THRESHOLDS.to(prediction.device):
predicted, observed = prediction >= threshold, target >= threshold
values.extend([(predicted & observed).sum(), (predicted & ~observed).sum(), (~predicted & observed).sum()])
return torch.stack(values).to(torch.float64)
def metrics_from_sums(sums: torch.Tensor) -> dict[str, float]:
count = sums[3].clamp_min(1)
csi = []
for index in range(len(CSI_THRESHOLDS)):
hits, false_alarms, misses = sums[4 + index * 3 : 7 + index * 3]
csi.append((hits / (hits + false_alarms + misses).clamp_min(1)).item())
return {
"mse": (sums[0] / count).item(),
"mae": (sums[1] / count).item(),
"ssim": (sums[2] / count).item(),
"mean_csi": sum(csi) / len(csi),
}
def metric_sums_light(prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
"""Additive MSE/MAE sums for cheap per-epoch validation: [mse_sum, mae_sum, count]."""
prediction, target = prediction.float(), target.float()
reduce_dims = tuple(range(1, prediction.ndim))
return torch.stack(
[
(prediction - target).square().mean(reduce_dims).sum(),
(prediction - target).abs().mean(reduce_dims).sum(),
prediction.new_tensor(prediction.shape[0]),
]
).to(torch.float64)
def metrics_from_light_sums(sums: torch.Tensor) -> dict[str, float]:
count = sums[2].clamp_min(1)
return {"mse": (sums[0] / count).item(), "mae": (sums[1] / count).item()}
def compute_metrics(prediction: torch.Tensor, target: torch.Tensor) -> dict[str, float]:
return metrics_from_sums(metric_sums(prediction.detach(), target.detach()))
|