Gaussian blur with kernelsize and sigma API

Author: tejank10

torchvision/transforms/functional.py

@@ -1028,18 +1028,18 @@ def erase(img: Tensor, i: int, j: int, h: int, w: int, v: Tensor, inplace: bool
     return img
 
 
-def gaussian_blur(img: Tensor, kernel_size: int, sigma: float = None) -> Tensor:
+def gaussian_blur(img: Tensor, kernel_size: List[int], sigma: Optional[List[float]] = None) -> Tensor:
     """Performs Gaussian blurring on the img by given kernel.
     The image can be a PIL Image or a Tensor, in which case it is expected
     to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions
 
     Args:
         img (PIL Image or Tensor): Image to be blurred
-        kernel_size (sequence or int): Gaussian kernel size. Can be a sequence of integers
+        kernel_size (sequence of ints or int): Gaussian kernel size. Can be a sequence of integers
             like ``(kx, ky)`` or a single integer for square kernels.
             In torchscript mode kernel_size as single int is not supported, use a tuple or
             list of length 1: ``[size, ]``.
-        sigma (sequence or float, optional): Gaussian kernel standard deviation. Can be a
+        sigma (sequence of floats or float or None, optional): Gaussian kernel standard deviation. Can be a
             sequence of floats like ``(sigma_x, sigma_y)`` or a single float to define the
             same sigma in both X/Y directions. If None, then it is computed using
             ``kernel_size`` as ``sigma = 0.3 * ((kernel_size - 1) * 0.5 - 1) + 0.8``.
@@ -1056,12 +1056,10 @@ def gaussian_blur(img: Tensor, kernel_size: int, sigma: float = None) -> Tensor:
             raise TypeError('img should be PIL Image or Tensor. Got {}'.format(type(img)))
 
         is_pil_image = True
-        t_img = pil_to_tensor(img)
+        t_img = to_tensor(img)
 
     output = F_t.gaussian_blur(t_img, kernel_size, sigma)
 
     if is_pil_image:
-        output = output.permute((1, 2, 0))
-        output = Image.fromarray(output.numpy())
-
+        output = to_pil_image(output)
     return output

torchvision/transforms/functional_tensor.py

@@ -942,65 +942,78 @@ def perspective(
     return _apply_grid_transform(img, grid, mode)
 
 
-def _get_kernel(radius: float, passes: int):
-    sigma2 = torch.Tensor([radius ** 2 / passes])
+def _get_gaussian_kernel1d(kernel_size: int, sigma: float):
+    ksize_half = (kernel_size - 1) * 0.5
 
-    kernel_rad = (torch.sqrt(12. * sigma2 + 1.) - 1.) / 2.
+    x = torch.linspace(-ksize_half, ksize_half, steps=kernel_size)
+    pdf = torch.exp(-0.5 * (x / sigma).pow(2))
+    kernel1d = pdf / pdf.sum()
 
-    kernel_rad_int = kernel_rad.long().item()
+    return kernel1d
 
-    kernel_rad_float = (2 * kernel_rad_int + 1) * (kernel_rad_int * (kernel_rad_int + 1) - 3 * sigma2)
-    kernel_rad_float /= 6 * (sigma2 - (kernel_rad_int + 1) * (kernel_rad_int + 1))
-    kernel_rad_float = kernel_rad_float.item()
 
-    kernel_rad = kernel_rad_int + kernel_rad_float
+def _get_gaussian_kernel2d(kernel_size: List[int], sigma: List[float]):
+    ksize_x, ksize_y = kernel_size
+    sigma_x, sigma_y = sigma
 
-    ksize = 2 * kernel_rad_int + 1 + 2 * (kernel_rad_float > 0)
-    kernel1d = torch.ones(ksize) / (2 * kernel_rad + 1)
+    kernel1d_x = _get_gaussian_kernel1d(ksize_x, sigma_x)
+    kernel1d_y = _get_gaussian_kernel1d(ksize_y, sigma_y)
 
-    if kernel_rad_float > 0:
-        kernel1d[[0, -1]] = kernel_rad_float / (2 * kernel_rad + 1)
-
-    kernel2d = torch.mm(kernel1d[:, None], kernel1d[None, :])
+    kernel2d = torch.mm(kernel1d_y[:, None], kernel1d_x[None, :])
 
     return kernel2d
 
 
-def gaussian_blur(img: Tensor, radius: float) -> Tensor:
+def gaussian_blur(img: Tensor, kernel_size: List[int], sigma: Optional[List[float]] = None) -> Tensor:
     """Performs Gaussian blurring on the img by given kernel.
 
     Args:
         img (Tensor): Image to be blurred
-        radius (float): Blur radius
+        kernel_size (sequence of int or int): Kernel size of the Gaussian kernel
+        sigma (sequence of float or float or None): Standard deviation of the Gaussian kernel
 
     Returns:
-        Tensor: An image that is blurred using kernel of given radius
+        Tensor: An image that is blurred using gaussian kernel of given parameters
     """
-    if not (isinstance(img, torch.Tensor) and _is_tensor_a_torch_image(img)):
+    if not (isinstance(img, torch.Tensor) or _is_tensor_a_torch_image(img)):
         raise TypeError('img should be Tensor Image. Got {}'.format(type(img)))
-    if not isinstance(radius, (float, int)):
-        raise TypeError('radius should be either float or int. Got {}'.format(type(radius)))
+    if not isinstance(kernel_size, (int, list, tuple)):
+        raise TypeError('kernel_size should be int or a sequence of integers. Got {}'.format(type(kernel_size)))
+    if not isinstance(sigma, (float, int, list, tuple)) and sigma != None:
+        raise TypeError('sigma should be either float or int or its sequence. Got {}'.format(type(sigma)))
+
+    if isinstance(kernel_size, int):
+        kernel_size = [kernel_size] * 2
+    if isinstance(sigma, (int, float, None)):
+        sigma = [sigma] * 2
+
+    if len(kernel_size) != 2:
+        raise ValueError('If kernel_size is a sequence its length should be 2. Got {}'.format(len(kernel_size)))
+    if len(sigma) != 2:
+        raise ValueError('If sigma is a sequence its length should be 2. Got {}'.format(len(sigma)))
+
+    if any([ksize % 2 == 0 or not isinstance(ksize, int) for ksize in kernel_size]):
+        raise ValueError('kernel_size should have odd and positive integers. Got {}'.format(kernel_size))
+
+    sigma =  [s if s != None else 0.3 * ((ksize - 1) * 0.5 - 1) + 0.8 for ksize, s in zip(kernel_size, sigma)]
 
-    radius = float(radius)
-    passes = 3
+    if any([s <= 0. for s in sigma]):
+        raise ValueError('sigma should have positive values. Got {}'.format(sigma))
 
     ndim = img.ndim
     if ndim == 2:
         img = img.unsqueeze(0)
     if ndim == 3:
         img = img.unsqueeze(0)
 
-    kernel = _get_kernel(radius, passes)
+    kernel = _get_gaussian_kernel2d(kernel_size, sigma)
 
-    padding = _compute_padding(kernel.shape[::-1])
+    padding = _compute_padding(kernel_size)
 
     kernel = kernel[None, None, :, :].repeat(img.size(-3), 1, 1, 1)
 
-    padded_img = pad(img, padding, padding_mode='edge')
+    padded_img = pad(img, padding, padding_mode='reflect')
     blurred_img = conv2d(padded_img, kernel, groups=img.size(-3))
-    for _ in range(passes - 1):
-        padded_img = pad(blurred_img, padding, padding_mode='edge')
-        blurred_img = conv2d(padded_img, kernel, groups=img.size(-3))
 
     if ndim == 2:
         return blurred_img[0, 0]

torchvision/transforms/transforms.py

@@ -1553,44 +1553,52 @@ class GaussianBlur(torch.nn.Module):
     dimensions
 
     Args:
-        radius (float or tuple of float (min, max)): Radius to be used for creating
-        kernel to perform blurring. If float, radius is fixed. If it is tuple of
-        float (min, max), kernel radius is chosen uniformly at random to lie in the
+        ksize (int): Size of the Gaussian kernel.
+        sigma (float or tuple of float (min, max)): Standard deviation to be used for
+        creating kernel to perform blurring. If float, sigma is fixed. If it is tuple
+        of float (min, max), sigma is chosen uniformly at random to lie in the
         given range.
 
     Returns:
         PIL Image or Tensor: Gaussian blurred version of the input image.
 
     """
 
-    def __init__(self, radius=(0.1, 2.0)):
+    def __init__(self, ksize, sigma=(0.1, 2.0)):
         super().__init__()
 
-        if isinstance(radius, numbers.Number):
-            if radius <= 0:
-                raise ValueError("If radius is a single number, it must be positive.")
-            radius = (radius, radius)
-        elif isinstance(radius, (tuple, list)) and len(radius) == 2:
-            if not 0. < radius[0] <= radius[1]:
-                raise ValueError("radius values should be positive and of the form (min, max)")
+        if isinstance(ksize, numbers.Number):
+            if ksize <= 0 or ksize % 2 == 0:
+                raise ValueError("ksize should be an odd and positive number.")
         else:
-            raise TypeError("radius should be a single number or a list/tuple with length 2.")
+            raise TypeError("ksize should be a single number.")
+
+        if isinstance(sigma, numbers.Number):
+            if sigma <= 0:
+                raise ValueError("If sigma is a single number, it must be positive.")
+            sigma = (sigma, sigma)
+        elif isinstance(sigma, (tuple, list)) and len(sigma) == 2:
+            if not 0. < sigma[0] <= sigma[1]:
+                raise ValueError("sigma values should be positive and of the form (min, max).")
+        else:
+            raise TypeError("sigma should be a single number or a list/tuple with length 2.")
 
-        self.rad_min, self.rad_max = radius
+        self.ksize = ksize
+        self.sigma_min, self.sigma_max = sigma
 
     @staticmethod
-    def get_params(rad_min: float, rad_max: float):
-        """Choose radius for ``gaussian_blur`` for random gaussian blurring.
+    def get_params(sigma_min: float, sigma_max: float):
+        """Choose sigma for ``gaussian_blur`` for random gaussian blurring.
 
         Args:
-            rad_min (float): Minimum radius that can be chosen for blurring kernel.
-            rad_max (float): Maximum radius that can be chosen for blurring kernel.
+            sigma_min (float): Minimum standard deviation that can be chosen for blurring kernel.
+            sigma_max (float): Maximum standard deviation that can be chosen for blurring kernel.
 
         Returns:
-            float: radius be passed to ``gaussian_blur`` for gaussian blurring.
+            float: Standard deviation to be passed to calculate kernel for gaussian blurring.
         """
-        radius = random.uniform(rad_min, rad_max)
-        return radius
+        sigma = random.uniform(sigma_min, sigma_max)
+        return sigma
 
     def forward(self, img):
         """
@@ -1600,8 +1608,10 @@ def forward(self, img):
         Returns:
             PIL Image or Tensor: Gaussian blurred image
         """
-        radius = self.get_params(self.rad_min, self.rad_max)
-        return F.gaussian_blur(img, radius)
+        sigma = self.get_params(self.sigma_min, self.sigma_max)
+        return F.gaussian_blur(img, self.ksize, sigma)
 
     def __repr__(self):
-        return self.__class__.__name__ + '(rad_min={0}, rad_max={1})'.format(self.rad_min, self.rad_max)
+        s = 'kernel size={0}, '.format(self.ksize)
+        s += '(sigma_min={0}, sigma_max={1})'.format(self.sigma_min, self.sigma_max)
+        return self.__class__.__name__ + s