| import torchvision.transforms.functional as TF | |
| def img_tensor_mae(tensor1, tensor2): | |
| """Calculate the mean absolute difference between two image tensors.""" | |
| # Remove batch dimensions if present | |
| tensor1 = tensor1.squeeze(0).cpu() | |
| tensor2 = tensor2.squeeze(0).cpu() | |
| if tensor1.shape != tensor2.shape: | |
| raise ValueError( | |
| f"Tensors must have the same shape for comparison. Got {tensor1.shape=} and {tensor2.shape=}." | |
| ) | |
| return (tensor1 - tensor2).abs().mean().item() | |
| def blur(tensor, kernel_size=9, sigma=None): | |
| """Apply Gaussian blur to an image tensor.""" | |
| # [1, H, W, C] -> [1, C, H, W] | |
| if tensor.ndim == 4: | |
| tensor = tensor.permute(0, 3, 1, 2) | |
| elif tensor.ndim == 3: | |
| tensor = tensor.permute(2, 0, 1).unsqueeze(0) | |
| else: | |
| raise ValueError(f"Expected a 3D or 4D tensor, got {tensor.ndim=}") | |
| return TF.gaussian_blur(tensor, kernel_size=kernel_size, sigma=sigma).permute( # type: ignore | |
| 0, 2, 3, 1 | |
| ) | |