Spaces:
Running
on
A10G
Running
on
A10G
| from typing import Tuple | |
| import torch | |
| import torch.nn as nn | |
| from .filter import filter2d, filter2d_separable | |
| from .kernels import get_gaussian_kernel1d, get_gaussian_kernel2d | |
| def gaussian_blur2d( | |
| input: torch.Tensor, | |
| kernel_size: Tuple[int, int], | |
| sigma: Tuple[float, float], | |
| border_type: str = 'reflect', | |
| separable: bool = True, | |
| ) -> torch.Tensor: | |
| r"""Create an operator that blurs a tensor using a Gaussian filter. | |
| .. image:: _static/img/gaussian_blur2d.png | |
| The operator smooths the given tensor with a gaussian kernel by convolving | |
| it to each channel. It supports batched operation. | |
| Arguments: | |
| input: the input tensor with shape :math:`(B,C,H,W)`. | |
| kernel_size: the size of the kernel. | |
| sigma: the standard deviation of the kernel. | |
| border_type: the padding mode to be applied before convolving. | |
| The expected modes are: ``'constant'``, ``'reflect'``, | |
| ``'replicate'`` or ``'circular'``. Default: ``'reflect'``. | |
| separable: run as composition of two 1d-convolutions. | |
| Returns: | |
| the blurred tensor with shape :math:`(B, C, H, W)`. | |
| .. note:: | |
| See a working example `here <https://kornia-tutorials.readthedocs.io/en/latest/ | |
| gaussian_blur.html>`__. | |
| Examples: | |
| >>> input = torch.rand(2, 4, 5, 5) | |
| >>> output = gaussian_blur2d(input, (3, 3), (1.5, 1.5)) | |
| >>> output.shape | |
| torch.Size([2, 4, 5, 5]) | |
| """ | |
| if separable: | |
| kernel_x: torch.Tensor = get_gaussian_kernel1d(kernel_size[1], sigma[1]) | |
| kernel_y: torch.Tensor = get_gaussian_kernel1d(kernel_size[0], sigma[0]) | |
| out = filter2d_separable(input, kernel_x[None], kernel_y[None], border_type) | |
| else: | |
| kernel: torch.Tensor = get_gaussian_kernel2d(kernel_size, sigma) | |
| out = filter2d(input, kernel[None], border_type) | |
| return out | |
| class GaussianBlur2d(nn.Module): | |
| r"""Create an operator that blurs a tensor using a Gaussian filter. | |
| The operator smooths the given tensor with a gaussian kernel by convolving | |
| it to each channel. It supports batched operation. | |
| Arguments: | |
| kernel_size: the size of the kernel. | |
| sigma: the standard deviation of the kernel. | |
| border_type: the padding mode to be applied before convolving. | |
| The expected modes are: ``'constant'``, ``'reflect'``, | |
| ``'replicate'`` or ``'circular'``. Default: ``'reflect'``. | |
| separable: run as composition of two 1d-convolutions. | |
| Returns: | |
| the blurred tensor. | |
| Shape: | |
| - Input: :math:`(B, C, H, W)` | |
| - Output: :math:`(B, C, H, W)` | |
| Examples:: | |
| >>> input = torch.rand(2, 4, 5, 5) | |
| >>> gauss = GaussianBlur2d((3, 3), (1.5, 1.5)) | |
| >>> output = gauss(input) # 2x4x5x5 | |
| >>> output.shape | |
| torch.Size([2, 4, 5, 5]) | |
| """ | |
| def __init__( | |
| self, | |
| kernel_size: Tuple[int, int], | |
| sigma: Tuple[float, float], | |
| border_type: str = 'reflect', | |
| separable: bool = True, | |
| ) -> None: | |
| super().__init__() | |
| self.kernel_size: Tuple[int, int] = kernel_size | |
| self.sigma: Tuple[float, float] = sigma | |
| self.border_type = border_type | |
| self.separable = separable | |
| def __repr__(self) -> str: | |
| return ( | |
| self.__class__.__name__ | |
| + '(kernel_size=' | |
| + str(self.kernel_size) | |
| + ', ' | |
| + 'sigma=' | |
| + str(self.sigma) | |
| + ', ' | |
| + 'border_type=' | |
| + self.border_type | |
| + 'separable=' | |
| + str(self.separable) | |
| + ')' | |
| ) | |
| def forward(self, input: torch.Tensor) -> torch.Tensor: | |
| return gaussian_blur2d(input, self.kernel_size, self.sigma, self.border_type, self.separable) |