Merge branch 'feat/adjust_hue_cvcuda' into feat/brightness_contrast_cvcuda

Author: justincdavis

test/test_transforms_v2.py

@@ -6149,7 +6149,18 @@ def test_kernel_image(self, dtype, device):
     def test_kernel_video(self):
         check_kernel(F.adjust_hue_video, make_video(), hue_factor=0.25)
 
-    @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_hue, make_input(), hue_factor=0.25)
 
@@ -6160,9 +6171,16 @@ def test_functional(self, make_input):
             (F._color._adjust_hue_image_pil, PIL.Image.Image),
             (F.adjust_hue_image, tv_tensors.Image),
             (F.adjust_hue_video, tv_tensors.Video),
+            pytest.param(
+                F._color._adjust_hue_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_hue, kernel=kernel, input_type=input_type)
 
     def test_functional_error(self):
@@ -6173,11 +6191,27 @@ def test_functional_error(self):
             with pytest.raises(ValueError, match=re.escape("is not in [-0.5, 0.5]")):
                 F.adjust_hue(make_image(), hue_factor=hue_factor)
 
+    @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("hue_factor", [-0.5, -0.3, 0.0, 0.2, 0.5])
-    def test_correctness_image(self, hue_factor):
-        image = make_image(dtype=torch.uint8, device="cpu")
+    def test_correctness_image(self, make_input, hue_factor):
+        image = make_input(dtype=torch.uint8, device="cpu")
 
         actual = F.adjust_hue(image, hue_factor=hue_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_hue(F.to_pil_image(image), hue_factor=hue_factor))
 
         mae = (actual.float() - expected.float()).abs().mean()

torchvision/transforms/v2/functional/_color.py

@@ -457,6 +457,31 @@ def adjust_hue_video(video: torch.Tensor, hue_factor: float) -> torch.Tensor:
     return adjust_hue_image(video, hue_factor=hue_factor)
 
 
+def _adjust_hue_cvcuda(image: "cvcuda.Tensor", hue_factor: float) -> "cvcuda.Tensor":
+    cvcuda = _import_cvcuda()
+
+    if not (-0.5 <= hue_factor <= 0.5):
+        raise ValueError(f"hue_factor ({hue_factor}) is not in [-0.5, 0.5].")
+
+    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
+
+    # no native adjust_hue, use CV-CUDA for color converison, use torch for elementwise operations
+    hsv = cvcuda.cvtcolor(image, cvcuda.ColorConversion.RGB2HSV)
+    hsv_torch = torch.as_tensor(hsv.cuda()).float()
+    hsv_torch[..., 0] = (hsv_torch[..., 0] + hue_factor * 180) % 180
+    hsv_modified = cvcuda.as_tensor(hsv_torch.to(torch.uint8), "NHWC")
+    return cvcuda.cvtcolor(hsv_modified, cvcuda.ColorConversion.HSV2RGB)
+
+
+if CVCUDA_AVAILABLE:
+    _register_kernel_internal(adjust_hue, _import_cvcuda().Tensor)(_adjust_hue_cvcuda)
+
+
 def adjust_gamma(inpt: torch.Tensor, gamma: float, gain: float = 1) -> torch.Tensor:
     """Adjust gamma."""
     if torch.jit.is_scripting():