Spaces:
Running
on
A10G
Running
on
A10G
File size: 3,815 Bytes
320e465 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
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) |