Merge branch 'feat/adjust_saturation_cvcuda' into brightness_contrast_hue_cvcuda

Author: justincdavis

test/test_transforms_v2.py

@@ -6225,7 +6225,18 @@ def test_kernel_image(self, dtype, device):
     def test_kernel_video(self):
         check_kernel(F.adjust_saturation_video, make_video(), saturation_factor=0.5)
 
-    @pytest.mark.parametrize("make_input", [make_image_tensor, make_image, make_image_pil, make_video])
+    @pytest.mark.parametrize(
+        "make_input",
+        [
+            make_image_tensor,
+            make_image,
+            make_image_pil,
+            make_video,
+            pytest.param(
+                make_image_cvcuda, marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="CVCUDA not available")
+            ),
+        ],
+    )
     def test_functional(self, make_input):
         check_functional(F.adjust_saturation, make_input(), saturation_factor=0.5)
 
@@ -6236,9 +6247,16 @@ def test_functional(self, make_input):
             (F._color._adjust_saturation_image_pil, PIL.Image.Image),
             (F.adjust_saturation_image, tv_tensors.Image),
             (F.adjust_saturation_video, tv_tensors.Video),
+            pytest.param(
+                F._color._adjust_saturation_cvcuda,
+                "cvcuda.Tensor",
+                marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="CVCUDA not available"),
+            ),
         ],
     )
     def test_functional_signature(self, kernel, input_type):
+        if input_type == "cvcuda.Tensor":
+            input_type = _import_cvcuda().Tensor
         check_functional_kernel_signature_match(F.adjust_saturation, kernel=kernel, input_type=input_type)
 
     def test_functional_error(self):
@@ -6248,11 +6266,28 @@ def test_functional_error(self):
         with pytest.raises(ValueError, match="is not non-negative"):
             F.adjust_saturation(make_image(), saturation_factor=-1)
 
+    @pytest.mark.parametrize(
+        "make_input",
+        [
+            make_image,
+            pytest.param(
+                make_image_cvcuda, marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="CVCUDA not available")
+            ),
+        ],
+    )
+    @pytest.mark.parametrize("color_space", ["RGB", "GRAY"])
     @pytest.mark.parametrize("saturation_factor", [0.1, 0.5, 1.0])
-    def test_correctness_image(self, saturation_factor):
-        image = make_image(dtype=torch.uint8, device="cpu")
+    def test_correctness_image(self, make_input, color_space, saturation_factor):
+        image = make_input(dtype=torch.uint8, color_space=color_space, device="cpu")
 
         actual = F.adjust_saturation(image, saturation_factor=saturation_factor)
+
+        if make_input is make_image_cvcuda:
+            actual = F.cvcuda_to_tensor(actual).to(device="cpu")
+            actual = actual.squeeze(0)
+            image = F.cvcuda_to_tensor(image)
+            image = image.squeeze(0)
+
         expected = F.to_image(F.adjust_saturation(F.to_pil_image(image), saturation_factor=saturation_factor))
 
         assert_close(actual, expected, rtol=0, atol=1)

torchvision/transforms/v2/functional/_color.py

@@ -197,6 +197,37 @@ def adjust_saturation_video(video: torch.Tensor, saturation_factor: float) -> to
     return adjust_saturation_image(video, saturation_factor=saturation_factor)
 
 
+def _adjust_saturation_cvcuda(image: "cvcuda.Tensor", saturation_factor: float) -> "cvcuda.Tensor":
+    if saturation_factor < 0:
+        raise ValueError(f"saturation_factor ({saturation_factor}) is not non-negative.")
+
+    c = image.shape[3]
+    if c not in [1, 3, 4]:
+        raise TypeError(f"Input image tensor permitted channel values are 1, 3, or 4, but found {c}")
+
+    if c == 1:  # Match PIL behaviour
+        return image
+
+    # grayscale weights
+    sf = saturation_factor
+    r, g, b = 0.2989, 0.587, 0.114
+    twist_data = [
+        [sf + (1 - sf) * r, (1 - sf) * g, (1 - sf) * b, 0.0],
+        [(1 - sf) * r, sf + (1 - sf) * g, (1 - sf) * b, 0.0],
+        [(1 - sf) * r, (1 - sf) * g, sf + (1 - sf) * b, 0.0],
+    ]
+    twist_tensor = cvcuda.as_tensor(
+        torch.tensor(twist_data, dtype=torch.float32, device="cuda"),
+        "HW",
+    )
+
+    return cvcuda.color_twist(image, twist_tensor)
+
+
+if CVCUDA_AVAILABLE:
+    _register_kernel_internal(adjust_saturation, cvcuda.Tensor)(_adjust_saturation_cvcuda)
+
+
 def adjust_contrast(inpt: torch.Tensor, contrast_factor: float) -> torch.Tensor:
     """See :class:`~torchvision.transforms.RandomAutocontrast`"""
     if torch.jit.is_scripting():