Added GaussianBlur transform and tests

Author: vfdev-5

test/test_prototype_transforms.py

@@ -644,3 +644,59 @@ def test_forward(self, padding, pad_if_needed, fill, padding_mode, mocker):
         else:
             # vfdev-5: I do not know how to mock and test this case
             pass
+
+
+class TestGaussianBlur:
+    def test_assertions(self):
+        with pytest.raises(ValueError, match="Kernel size should be a tuple/list of two integers"):
+            transforms.GaussianBlur([10, 12, 14])
+
+        with pytest.raises(ValueError, match="Kernel size value should be an odd and positive number"):
+            transforms.GaussianBlur(4)
+
+        with pytest.raises(TypeError, match="sigma should be a single float or a list/tuple with length 2"):
+            transforms.GaussianBlur(3, sigma=[1, 2, 3])
+
+        with pytest.raises(ValueError, match="If sigma is a single number, it must be positive"):
+            transforms.GaussianBlur(3, sigma=-1.0)
+
+        with pytest.raises(ValueError, match="sigma values should be positive and of the form"):
+            transforms.GaussianBlur(3, sigma=[2.0, 1.0])
+
+    @pytest.mark.parametrize("sigma", [10.0, [10.0, 12.0]])
+    def test__get_params(self, sigma):
+        transform = transforms.GaussianBlur(3, sigma=sigma)
+        params = transform._get_params(None)
+
+        if isinstance(sigma, float):
+            assert params["sigma"][0] == params["sigma"][1] == 10
+        else:
+            assert sigma[0] <= params["sigma"][0] <= sigma[1]
+            assert sigma[0] <= params["sigma"][1] <= sigma[1]
+
+    @pytest.mark.parametrize("kernel_size", [3, [3, 5], (5, 3)])
+    @pytest.mark.parametrize("sigma", [2.0, [2.0, 3.0]])
+    def test__transform(self, kernel_size, sigma, mocker):
+        transform = transforms.GaussianBlur(kernel_size=kernel_size, sigma=sigma)
+
+        if isinstance(kernel_size, (tuple, list)):
+            assert transform.kernel_size == kernel_size
+        else:
+            assert transform.kernel_size == (kernel_size, kernel_size)
+
+        if isinstance(sigma, (tuple, list)):
+            assert transform.sigma == sigma
+        else:
+            assert transform.sigma == (sigma, sigma)
+
+
+        fn = mocker.patch("torchvision.prototype.transforms.functional.gaussian_blur")
+        inpt = features.Image(torch.rand(1, 3, 32, 32))
+        # vfdev-5, Feature Request: let's store params as Transform attribute
+        # This could be also helpful for users
+        torch.manual_seed(12)
+        _ = transform(inpt)
+        torch.manual_seed(12)
+        params = transform._get_params(inpt)
+
+        fn.assert_called_once_with(inpt, **params)

test/test_prototype_transforms_functional.py

@@ -495,6 +495,7 @@ def center_crop_bounding_box():
         )
 
 
+@register_kernel_info_from_sample_inputs_fn
 def center_crop_segmentation_mask():
     for mask, output_size in itertools.product(
         make_segmentation_masks(image_sizes=((16, 16), (7, 33), (31, 9))),
@@ -503,6 +504,16 @@ def center_crop_segmentation_mask():
         yield SampleInput(mask, output_size)
 
 
+@register_kernel_info_from_sample_inputs_fn
+def gaussian_blur_image_tensor():
+    for image, kernel_size, sigma in itertools.product(
+        make_images(extra_dims=((4,),)),
+        [[3, 3], ],
+        [None, [3.0, 3.0]],
+    ):
+        yield SampleInput(image, kernel_size=kernel_size, sigma=sigma)
+
+
 @pytest.mark.parametrize(
     "kernel",
     [

torchvision/prototype/features/_feature.py

@@ -189,3 +189,6 @@ def equalize(self) -> Any:
 
     def invert(self) -> Any:
         return self
+
+    def gaussian_blur(self, kernel_size: List[int], sigma: Optional[List[float]] = None) -> Any:
+        return self

torchvision/prototype/features/_image.py

@@ -309,3 +309,9 @@ def invert(self) -> Image:
 
         output = _F.invert_image_tensor(self)
         return Image.new_like(self, output)
+
+    def gaussian_blur(self, kernel_size: List[int], sigma: Optional[List[float]] = None) -> Image:
+        from torchvision.prototype.transforms import functional as _F
+
+        output = _F.gaussian_blur_image_tensor(self, kernel_size=kernel_size, sigma=sigma)
+        return Image.new_like(self, output)

torchvision/prototype/transforms/__init__.py

@@ -22,7 +22,10 @@
     RandomAffine,
 )
 from ._meta import ConvertBoundingBoxFormat, ConvertImageDtype, ConvertImageColorSpace
-from ._misc import Identity, Normalize, ToDtype, Lambda
+from ._misc import Identity, GaussianBlur, Normalize, ToDtype, Lambda
 from ._type_conversion import DecodeImage, LabelToOneHot
 
 from ._deprecated import Grayscale, RandomGrayscale, ToTensor, ToPILImage, PILToTensor  # usort: skip
+
+# TODO: add RandomPerspective, RandomInvert, RandomPosterize, RandomSolarize,
+# RandomAdjustSharpness, RandomAutocontrast, ElasticTransform

torchvision/prototype/transforms/_misc.py

@@ -1,7 +1,8 @@
 import functools
-from typing import Any, List, Type, Callable, Dict
+from typing import Any, List, Type, Callable, Dict, Sequence, Union
 
 import torch
+from torchvision.transforms.transforms import _setup_size
 from torchvision.prototype.transforms import Transform, functional as F
 
 
@@ -46,6 +47,36 @@ def _transform(self, input: Any, params: Dict[str, Any]) -> Any:
             return input
 
 
+class GaussianBlur(Transform):
+    def __init__(
+        self, kernel_size: Union[int, Sequence[int]], sigma: Union[float, Sequence[float]] = (0.1, 2.0)
+    ) -> None:
+        super().__init__()
+        self.kernel_size = _setup_size(kernel_size, "Kernel size should be a tuple/list of two integers")
+        for ks in self.kernel_size:
+            if ks <= 0 or ks % 2 == 0:
+                raise ValueError("Kernel size value should be an odd and positive number.")
+
+        if isinstance(sigma, float):
+            if sigma <= 0:
+                raise ValueError("If sigma is a single number, it must be positive.")
+            sigma = (sigma, sigma)
+        elif isinstance(sigma, Sequence) and len(sigma) == 2:
+            if not 0.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 float or a list/tuple with length 2 floats.")
+
+        self.sigma = sigma
+
+    def _get_params(self, sample: Any) -> Dict[str, Any]:
+        sigma = torch.empty(1).uniform_(self.sigma[0], self.sigma[1]).item()
+        return dict(sigma=[sigma, sigma])
+
+    def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        return F.gaussian_blur(inpt, **params)
+
+
 class ToDtype(Lambda):
     def __init__(self, dtype: torch.dtype, *types: Type) -> None:
         self.dtype = dtype

torchvision/prototype/transforms/functional/__init__.py

@@ -99,7 +99,12 @@
     ten_crop_image_tensor,
     ten_crop_image_pil,
 )
-from ._misc import normalize_image_tensor, gaussian_blur_image_tensor
+from ._misc import (
+    normalize_image_tensor,
+    gaussian_blur,
+    gaussian_blur_image_tensor,
+    gaussian_blur_image_pil,
+)
 from ._type_conversion import (
     decode_image_with_pil,
     decode_video_with_av,

torchvision/prototype/transforms/functional/_misc.py

@@ -1,14 +1,28 @@
-from typing import Optional, List
+from typing import Optional, List, Union
 
 import PIL.Image
 import torch
+from torchvision.prototype import features
 from torchvision.transforms import functional_tensor as _FT
 from torchvision.transforms.functional import pil_to_tensor, to_pil_image
 
 
+# shortcut type
+DType = Union[torch.Tensor, PIL.Image.Image, features._Feature]
+
+
 normalize_image_tensor = _FT.normalize
 
 
+def normalize(inpt: DType, mean: List[float], std: List[float], inplace: bool = False) -> DType:
+    if isinstance(inpt, features.Image):
+        return normalize_image_tensor(inpt, mean=mean, std=std, inplace=inplace)
+    elif type(inpt) == torch.Tensor:
+        return normalize_image_tensor(inpt, mean=mean, std=std, inplace=inplace)
+    else:
+        raise TypeError("Unsupported input type")
+
+
 def gaussian_blur_image_tensor(
     img: torch.Tensor, kernel_size: List[int], sigma: Optional[List[float]] = None
 ) -> torch.Tensor:
@@ -42,3 +56,12 @@ def gaussian_blur_image_pil(img: PIL.Image, kernel_size: List[int], sigma: Optio
     t_img = pil_to_tensor(img)
     output = gaussian_blur_image_tensor(t_img, kernel_size=kernel_size, sigma=sigma)
     return to_pil_image(output, mode=img.mode)
+
+
+def gaussian_blur(inpt: DType, kernel_size: List[int], sigma: Optional[List[float]] = None) -> DType:
+    if isinstance(inpt, features._Feature):
+        return inpt.gaussian_blur(kernel_size=kernel_size, sigma=sigma)
+    elif isinstance(inpt, PIL.Image.Image):
+        return gaussian_blur_image_pil(inpt, kernel_size=kernel_size, sigma=sigma)
+    else:
+        return gaussian_blur_image_tensor(inpt, kernel_size=kernel_size, sigma=sigma)