(
ground_truth: Float[Tensor, "batch channel height width"],
predicted: Float[Tensor, "batch channel height width"],
)
| 35 | |
| 36 | @torch.no_grad() |
| 37 | def compute_ssim( |
| 38 | ground_truth: Float[Tensor, "batch channel height width"], |
| 39 | predicted: Float[Tensor, "batch channel height width"], |
| 40 | ) -> Float[Tensor, " batch"]: |
| 41 | ssim = [ |
| 42 | structural_similarity( |
| 43 | gt.detach().cpu().numpy(), |
| 44 | hat.detach().cpu().numpy(), |
| 45 | win_size=11, |
| 46 | gaussian_weights=True, |
| 47 | channel_axis=0, |
| 48 | data_range=1.0, |
| 49 | ) |
| 50 | for gt, hat in zip(ground_truth, predicted) |
| 51 | ] |
| 52 | return torch.tensor(ssim, dtype=predicted.dtype, device=predicted.device) |
no outgoing calls
no test coverage detected