begin editing for main changes

Author: justincdavis

test/test_transforms_v2.py

@@ -25,7 +25,6 @@
     assert_equal,
     cache,
     cpu_and_cuda,
-    cvcuda_to_pil_compatible_tensor,
     freeze_rng_state,
     ignore_jit_no_profile_information_warning,
     make_bounding_boxes,
@@ -5848,14 +5847,14 @@ def test_functional(self, make_input):
             (F.invert_image, tv_tensors.Image),
             (F.invert_video, tv_tensors.Video),
             pytest.param(
-                F._color._invert_cvcuda,
-                "cvcuda.Tensor",
+                F._color._invert_image_cvcuda,
+                None,
                 marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA"),
             ),
         ],
     )
     def test_functional_signature(self, kernel, input_type):
-        if input_type == "cvcuda.Tensor":
+        if kernel is F._color._invert_image_cvcuda:
             input_type = _import_cvcuda().Tensor
         check_functional_kernel_signature_match(F.invert, kernel=kernel, input_type=input_type)
 
@@ -5890,7 +5889,7 @@ def test_correctness_image(self, make_input, fn):
         actual = fn(image)
 
         if make_input is make_image_cvcuda:
-            image = cvcuda_to_pil_compatible_tensor(image)
+            image = F.cvcuda_to_tensor(image)[0].cpu()
 
         expected = F.to_image(F.invert(F.to_pil_image(image)))
 

torchvision/transforms/v2/_color.py

@@ -5,11 +5,15 @@
 import torch
 from torchvision import transforms as _transforms
 from torchvision.transforms.v2 import functional as F, Transform
+from torchvision.transforms.v2.functional._utils import _is_cvcuda_available, _is_cvcuda_tensor
 
 from ._transform import _RandomApplyTransform
 from ._utils import query_chw
 
 
+CVCUDA_AVAILABLE = _is_cvcuda_available()
+
+
 class Grayscale(Transform):
     """Convert images or videos to grayscale.
 
@@ -282,6 +286,9 @@ class RandomInvert(_RandomApplyTransform):
 
     _v1_transform_cls = _transforms.RandomInvert
 
+    if CVCUDA_AVAILABLE:
+        _transformed_types = _RandomApplyTransform._transformed_types + (_is_cvcuda_tensor,)
+
     def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
         return self._call_kernel(F.invert, inpt)
 

torchvision/transforms/v2/functional/_color.py

@@ -690,11 +690,10 @@ def invert_video(video: torch.Tensor) -> torch.Tensor:
     return invert_image(video)
 
 
-if CVCUDA_AVAILABLE:
-    _invert_cvcuda_tensors = {}
+_invert_cvcuda_tensors = {}
 
 
-def _invert_cvcuda(image: "cvcuda.Tensor") -> "cvcuda.Tensor":
+def _invert_image_cvcuda(image: "cvcuda.Tensor") -> "cvcuda.Tensor":
     cvcuda = _import_cvcuda()
 
     # save the tensors into a dictionary only if CV-CUDA is actually used
@@ -725,7 +724,7 @@ def _invert_cvcuda(image: "cvcuda.Tensor") -> "cvcuda.Tensor":
 
 
 if CVCUDA_AVAILABLE:
-    _register_kernel_internal(invert, _import_cvcuda().Tensor)(_invert_cvcuda)
+    _register_kernel_internal(invert, _import_cvcuda().Tensor)(_invert_image_cvcuda)
 
 
 def permute_channels(inpt: torch.Tensor, permutation: list[int]) -> torch.Tensor: