adjust saturation complete and tested

Author: justincdavis

test/test_transforms_v2.py

@@ -6277,7 +6277,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)
 
@@ -6288,9 +6299,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):
@@ -6300,11 +6318,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

@@ -201,35 +201,31 @@ def _adjust_saturation_cvcuda(image: "cvcuda.Tensor", saturation_factor: float)
     if saturation_factor < 0:
         raise ValueError(f"saturation_factor ({saturation_factor}) is not non-negative.")
 
-    c = image.shape[-1]  # NHWC layout
+    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 (same as _rgb_to_grayscale_image)
+    # grayscale weights
     sf = saturation_factor
     r, g, b = 0.2989, 0.587, 0.114
-
-    # Build 3x4 saturation matrix
     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 = cvcuda.Tensor(
-        torch.tensor(twist_data, dtype=torch.float32, device="cuda").contiguous(),
-        layout="HW",
+    twist_tensor = cvcuda.as_tensor(
+        torch.tensor(twist_data, dtype=torch.float32, device="cuda"),
+        "HW",
     )
 
-    return cvcuda.color_twist(image, twist)
+    return cvcuda.color_twist(image, twist_tensor)
 
 
 if CVCUDA_AVAILABLE:
-    _register_kernel_internal(adjust_saturation, cvcuda.Tensor)(
-        _adjust_saturation_cvcuda
-    )
+    _register_kernel_internal(adjust_saturation, cvcuda.Tensor)(_adjust_saturation_cvcuda)
 
 
 def adjust_contrast(inpt: torch.Tensor, contrast_factor: float) -> torch.Tensor: