28 lines
1007 B
Python
28 lines
1007 B
Python
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
|
|
)
|