Spaces:
Running
Running
| # -*- coding: utf-8 -*- | |
| import torch | |
| import torch.nn.functional as F | |
| __all__ = ["filter2D", "gaussian", "gaussian_kernel2d", "sobel_hv"] | |
| def filter2D(input_tensor: torch.Tensor, kernel: torch.Tensor) -> torch.Tensor: | |
| """Convolves a given kernel on input tensor without losing dimensional shape. | |
| Parameters | |
| ---------- | |
| input_tensor : torch.Tensor | |
| Input image/tensor. | |
| kernel : torch.Tensor | |
| Convolution kernel/window. | |
| Returns | |
| ------- | |
| torch.Tensor: | |
| The convolved tensor of same shape as the input. | |
| """ | |
| (_, channel, _, _) = input_tensor.size() | |
| # "SAME" padding to avoid losing height and width | |
| pad = [ | |
| kernel.size(2) // 2, | |
| kernel.size(2) // 2, | |
| kernel.size(3) // 2, | |
| kernel.size(3) // 2, | |
| ] | |
| pad_tensor = F.pad(input_tensor, pad, "replicate") | |
| out = F.conv2d(pad_tensor, kernel, groups=channel) | |
| return out | |
| def gaussian( | |
| window_size: int, sigma: float, device: torch.device = None | |
| ) -> torch.Tensor: | |
| """Create a gaussian 1D tensor. | |
| Parameters | |
| ---------- | |
| window_size : int | |
| Number of elements for the output tensor. | |
| sigma : float | |
| Std of the gaussian distribution. | |
| device : torch.device | |
| Device for the tensor. | |
| Returns | |
| ------- | |
| torch.Tensor: | |
| A gaussian 1D tensor. Shape: (window_size, ). | |
| """ | |
| x = torch.arange(window_size, device=device).float() - window_size // 2 | |
| if window_size % 2 == 0: | |
| x = x + 0.5 | |
| gauss = torch.exp((-x.pow(2.0) / float(2 * sigma**2))) | |
| return gauss / gauss.sum() | |
| def gaussian_kernel2d( | |
| window_size: int, sigma: float, n_channels: int = 1, device: torch.device = None | |
| ) -> torch.Tensor: | |
| """Create 2D window_size**2 sized kernel a gaussial kernel. | |
| Parameters | |
| ---------- | |
| window_size : int | |
| Number of rows and columns for the output tensor. | |
| sigma : float | |
| Std of the gaussian distribution. | |
| n_channel : int | |
| Number of channels in the image that will be convolved with | |
| this kernel. | |
| device : torch.device | |
| Device for the kernel. | |
| Returns: | |
| ----------- | |
| torch.Tensor: | |
| A tensor of shape (1, 1, window_size, window_size) | |
| """ | |
| kernel_x = gaussian(window_size, sigma, device=device) | |
| kernel_y = gaussian(window_size, sigma, device=device) | |
| kernel_2d = torch.matmul(kernel_x.unsqueeze(-1), kernel_y.unsqueeze(-1).t()) | |
| kernel_2d = kernel_2d.expand(n_channels, 1, window_size, window_size) | |
| return kernel_2d | |
| def sobel_hv(window_size: int = 5, device: torch.device = None): | |
| """Create a kernel that is used to compute 1st order derivatives. | |
| Parameters | |
| ---------- | |
| window_size : int | |
| Size of the convolution kernel. | |
| device : torch.device: | |
| Device for the kernel. | |
| Returns | |
| ------- | |
| torch.Tensor: | |
| the computed 1st order derivatives of the input tensor. | |
| Shape (B, 2, H, W) | |
| Raises | |
| ------ | |
| ValueError: | |
| If `window_size` is not an odd number. | |
| """ | |
| if not window_size % 2 == 1: | |
| raise ValueError(f"window_size must be odd. Got: {window_size}") | |
| # Generate the sobel kernels | |
| range_h = torch.arange( | |
| -window_size // 2 + 1, window_size // 2 + 1, dtype=torch.float32, device=device | |
| ) | |
| range_v = torch.arange( | |
| -window_size // 2 + 1, window_size // 2 + 1, dtype=torch.float32, device=device | |
| ) | |
| h, v = torch.meshgrid(range_h, range_v) | |
| kernel_h = h / (h * h + v * v + 1e-6) | |
| kernel_h = kernel_h.unsqueeze(0).unsqueeze(0) | |
| kernel_v = v / (h * h + v * v + 1e-6) | |
| kernel_v = kernel_v.unsqueeze(0).unsqueeze(0) | |
| return torch.cat([kernel_h, kernel_v], dim=0) | |