draft initial gaussian_blur cvcuda kernel implementation

Author: justincdavis

test/test_transforms_v2.py

@@ -3995,6 +3995,70 @@ def test_functional_image_correctness(self, dimensions, kernel_size, sigma, dtyp
         torch.testing.assert_close(actual, expected, rtol=0, atol=1)
 
 
+@pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA")
+@needs_cuda
+class TestGaussianBlurCVCUDA:
+    def test_kernel_image_errors(self):
+        image = make_image_cvcuda(batch_dims=(1,))
+
+        with pytest.raises(ValueError, match="kernel_size is a sequence its length should be 2"):
+            F.gaussian_blur_cvcuda(image, kernel_size=[1, 2, 3])
+
+        for kernel_size in [2, -1]:
+            with pytest.raises(ValueError, match="kernel_size should have odd and positive integers"):
+                F.gaussian_blur_cvcuda(image, kernel_size=kernel_size)
+
+        with pytest.raises(ValueError, match="sigma is a sequence, its length should be 2"):
+            F.gaussian_blur_cvcuda(image, kernel_size=1, sigma=[1, 2, 3])
+
+        with pytest.raises(TypeError, match="sigma should be either float or sequence of floats"):
+            F.gaussian_blur_cvcuda(image, kernel_size=1, sigma=object())
+
+        with pytest.raises(ValueError, match="sigma should have positive values"):
+            F.gaussian_blur_cvcuda(image, kernel_size=1, sigma=-1)
+
+    def test_functional(self):
+        check_functional(F.gaussian_blur, make_image_cvcuda(batch_dims=(1,)), kernel_size=(3, 3))
+
+    @pytest.mark.parametrize("device", cpu_and_cuda())
+    @pytest.mark.parametrize("sigma", [5, 2.0, (0.5, 2), [1.3, 2.7]])
+    def test_transform(self, device, sigma):
+        check_transform(
+            transforms.GaussianBlur(kernel_size=3, sigma=sigma), make_image_cvcuda(batch_dims=(1,), device=device)
+        )
+
+    @pytest.mark.parametrize(
+        ("dimensions", "kernel_size", "sigma"),
+        [
+            ((10, 12), (3, 3), 0.8),
+            ((10, 12), (3, 3), 0.5),
+            ((10, 12), (3, 5), 0.8),
+            ((10, 12), (3, 5), 0.5),
+            ((26, 28), (23, 23), 1.7),
+        ],
+    )
+    @pytest.mark.parametrize("color_space", ["RGB", "GRAY"])
+    @pytest.mark.parametrize("batch_dims", [(1,), (2,), (4,)])
+    @pytest.mark.parametrize("dtype", [torch.uint8, torch.float32])
+    def test_functional_image_correctness(self, dimensions, kernel_size, sigma, color_space, batch_dims, dtype):
+        height, width = dimensions
+
+        image_tensor = make_image(
+            size=(height, width), color_space=color_space, batch_dims=batch_dims, dtype=dtype, device="cuda"
+        )
+        image_cvcuda = F.to_cvcuda_tensor(image_tensor)
+
+        expected = F.gaussian_blur_image(image_tensor, kernel_size=kernel_size, sigma=sigma)
+        actual = F.gaussian_blur_cvcuda(image_cvcuda, kernel_size=kernel_size, sigma=sigma)
+        actual_torch = F.cvcuda_to_tensor(actual)
+
+        if dtype.is_floating_point:
+            torch.testing.assert_close(actual_torch, expected, rtol=0, atol=0.3)
+        else:
+            # uint8/16 gaussians can differ by up to max-value, most likely an overflow issue
+            torch.testing.assert_close(actual_torch, expected, rtol=0, atol=get_max_value(dtype))
+
+
 class TestGaussianNoise:
     @pytest.mark.parametrize(
         "make_input",

torchvision/transforms/v2/functional/__init__.py

@@ -147,6 +147,7 @@
 from ._misc import (
     convert_image_dtype,
     gaussian_blur,
+    gaussian_blur_cvcuda,
     gaussian_blur_image,
     gaussian_blur_video,
     gaussian_noise,

torchvision/transforms/v2/functional/_misc.py

@@ -1,5 +1,5 @@
 import math
-from typing import Optional, TYPE_CHECKING
+from typing import Optional, Sequence, TYPE_CHECKING
 
 import PIL.Image
 import torch
@@ -106,11 +106,10 @@ def _get_gaussian_kernel2d(
     return kernel2d
 
 
-@_register_kernel_internal(gaussian_blur, torch.Tensor)
-@_register_kernel_internal(gaussian_blur, tv_tensors.Image)
-def gaussian_blur_image(
-    image: torch.Tensor, kernel_size: list[int], sigma: Optional[list[float]] = None
-) -> torch.Tensor:
+def _validate_kernel_size_and_sigma(
+    kernel_size: Sequence[int] | int,
+    sigma: Sequence[float | int] | float | int | None = None,
+) -> tuple[list[int], list[float]]:
     # TODO: consider deprecating integers from sigma on the future
     if isinstance(kernel_size, int):
         kernel_size = [kernel_size, kernel_size]
@@ -139,6 +138,16 @@ def gaussian_blur_image(
         if s <= 0.0:
             raise ValueError(f"sigma should have positive values. Got {sigma}")
 
+    return kernel_size, sigma
+
+
+@_register_kernel_internal(gaussian_blur, torch.Tensor)
+@_register_kernel_internal(gaussian_blur, tv_tensors.Image)
+def gaussian_blur_image(
+    image: torch.Tensor, kernel_size: list[int], sigma: Optional[list[float]] = None
+) -> torch.Tensor:
+    kernel_size, sigma = _validate_kernel_size_and_sigma(kernel_size, sigma)
+
     if image.numel() == 0:
         return image
 
@@ -188,6 +197,25 @@ def gaussian_blur_video(
     return gaussian_blur_image(video, kernel_size, sigma)
 
 
+def gaussian_blur_cvcuda(
+    image: "cvcuda.Tensor",
+    kernel_size: Sequence[int] | int,
+    sigma: Sequence[float | int] | float | int | None = None,
+) -> "cvcuda.Tensor":
+    kernel_size, sigma = _validate_kernel_size_and_sigma(kernel_size, sigma)
+
+    return cvcuda.gaussian(
+        image,
+        tuple(kernel_size),
+        tuple(sigma),
+        border=cvcuda.Border.REFLECT,
+    )
+
+
+if CVCUDA_AVAILABLE:
+    _gaussian_blur_cvcuda = _register_kernel_internal(gaussian_blur, cvcuda.Tensor)(gaussian_blur_cvcuda)
+
+
 def gaussian_noise(inpt: torch.Tensor, mean: float = 0.0, sigma: float = 0.1, clip: bool = True) -> torch.Tensor:
     """See :class:`~torchvision.transforms.v2.GaussianNoise`"""
     if torch.jit.is_scripting():