from collections.abc import Sequence from typing import Union import torch from torch import Tensor from torch.nn import functional as F # noqa: N812 def _gaussian(kernel_size: int, sigma: float, dtype: torch.dtype, device: Union[torch.device, str]) -> Tensor: """Compute 1D gaussian kernel. Args: kernel_size: size of the gaussian kernel sigma: Standard deviation of the gaussian kernel dtype: data type of the output tensor device: device of the output tensor Example: >>> _gaussian(3, 1, torch.float, 'cpu') tensor([[0.2741, 0.4519, 0.2741]]) """ dist = torch.arange(start=(1 - kernel_size) / 2, end=(1 + kernel_size) / 2, step=1, dtype=dtype, device=device) gauss = torch.exp(-torch.pow(dist / sigma, 2) / 2) return (gauss / gauss.sum()).unsqueeze(dim=0) # (1, kernel_size) def _gaussian_kernel_2d( channel: int, kernel_size: Sequence[int], sigma: Sequence[float], dtype: torch.dtype, device: Union[torch.device, str], ) -> Tensor: """Compute 2D gaussian kernel. Args: channel: number of channels in the image kernel_size: size of the gaussian kernel as a tuple (h, w) sigma: Standard deviation of the gaussian kernel dtype: data type of the output tensor device: device of the output tensor Example: >>> _gaussian_kernel_2d(1, (5,5), (1,1), torch.float, "cpu") tensor([[[[0.0030, 0.0133, 0.0219, 0.0133, 0.0030], [0.0133, 0.0596, 0.0983, 0.0596, 0.0133], [0.0219, 0.0983, 0.1621, 0.0983, 0.0219], [0.0133, 0.0596, 0.0983, 0.0596, 0.0133], [0.0030, 0.0133, 0.0219, 0.0133, 0.0030]]]]) """ gaussian_kernel_x = _gaussian(kernel_size[0], sigma[0], dtype, device) gaussian_kernel_y = _gaussian(kernel_size[1], sigma[1], dtype, device) kernel = torch.matmul(gaussian_kernel_x.t(), gaussian_kernel_y) # (kernel_size, 1) * (1, kernel_size) return kernel.expand(channel, 1, kernel_size[0], kernel_size[1]) def _uniform_weight_bias_conv2d(inputs: Tensor, window_size: int) -> tuple[Tensor, Tensor]: """Construct uniform weight and bias for a 2d convolution. Args: inputs: Input image window_size: size of convolutional kernel Return: The weight and bias for 2d convolution """ kernel_weight = torch.ones(1, 1, window_size, window_size, dtype=inputs.dtype, device=inputs.device) kernel_weight /= window_size**2 kernel_bias = torch.zeros(1, dtype=inputs.dtype, device=inputs.device) return kernel_weight, kernel_bias def _single_dimension_pad(inputs: Tensor, dim: int, pad: int, outer_pad: int = 0) -> Tensor: """Apply single-dimension reflection padding to match scipy implementation. Args: inputs: Input image dim: A dimension the image should be padded over pad: Number of pads outer_pad: Number of outer pads Return: Image padded over a single dimension """ _max = inputs.shape[dim] x = torch.index_select(inputs, dim, torch.arange(pad - 1, -1, -1).to(inputs.device)) y = torch.index_select(inputs, dim, torch.arange(_max - 1, _max - pad - outer_pad, -1).to(inputs.device)) return torch.cat((x, inputs, y), dim) def _reflection_pad_2d(inputs: Tensor, pad: int, outer_pad: int = 0) -> Tensor: """Apply reflection padding to the input image. Args: inputs: Input image pad: Number of pads outer_pad: Number of outer pads Return: Padded image """ for dim in [2, 3]: inputs = _single_dimension_pad(inputs, dim, pad, outer_pad) return inputs def _uniform_filter(inputs: Tensor, window_size: int) -> Tensor: """Apply uniform filter with a window of a given size over the input image. Args: inputs: Input image window_size: Sliding window used for rmse calculation Return: Image transformed with the uniform input """ inputs = _reflection_pad_2d(inputs, window_size // 2, window_size % 2) kernel_weight, kernel_bias = _uniform_weight_bias_conv2d(inputs, window_size) # Iterate over channels return torch.cat( [ F.conv2d(inputs[:, channel].unsqueeze(1), kernel_weight, kernel_bias, padding=0) for channel in range(inputs.shape[1]) ], dim=1, ) def _gaussian_kernel_3d( channel: int, kernel_size: Sequence[int], sigma: Sequence[float], dtype: torch.dtype, device: torch.device ) -> Tensor: """Compute 3D gaussian kernel. Args: channel: number of channels in the image kernel_size: size of the gaussian kernel as a tuple (h, w, d) sigma: Standard deviation of the gaussian kernel dtype: data type of the output tensor device: device of the output tensor """ gaussian_kernel_x = _gaussian(kernel_size[0], sigma[0], dtype, device) gaussian_kernel_y = _gaussian(kernel_size[1], sigma[1], dtype, device) gaussian_kernel_z = _gaussian(kernel_size[2], sigma[2], dtype, device) kernel_xy = torch.matmul(gaussian_kernel_x.t(), gaussian_kernel_y) # (kernel_size, 1) * (1, kernel_size) kernel = torch.mul( kernel_xy.unsqueeze(-1).repeat(1, 1, kernel_size[2]), gaussian_kernel_z.expand(kernel_size[0], kernel_size[1], kernel_size[2]), ) return kernel.expand(channel, 1, kernel_size[0], kernel_size[1], kernel_size[2]) def _reflection_pad_3d(inputs: Tensor, pad_h: int, pad_w: int, pad_d: int) -> Tensor: """Reflective padding of 3d input. Args: inputs: tensor to pad, should be a 3D tensor of shape ``[N, C, H, W, D]`` pad_w: amount of padding in the height dimension pad_h: amount of padding in the width dimension pad_d: amount of padding in the depth dimension Returns: padded input tensor """ return F.pad(inputs, (pad_h, pad_h, pad_w, pad_w, pad_d, pad_d), mode="reflect")