update to include five ten crop and resized crop, use placeholder transforms for flip and resize for now

Author: justincdavis

test/common_utils.py

@@ -20,19 +20,15 @@
 from torch.testing._comparison import BooleanPair, NonePair, not_close_error_metas, NumberPair, TensorLikePair
 from torchvision import io, tv_tensors
 from torchvision.transforms._functional_tensor import _max_value as get_max_value
-from torchvision.transforms.v2.functional import (
-    cvcuda_to_tensor,
-    is_cvcuda_tensor,
-    to_cvcuda_tensor,
-    to_image,
-    to_pil_image,
-)
+from torchvision.transforms.v2.functional import cvcuda_to_tensor, to_cvcuda_tensor, to_image, to_pil_image
+from torchvision.transforms.v2.functional._utils import _import_cvcuda, _is_cvcuda_available
 from torchvision.utils import _Image_fromarray
 
 
 IN_OSS_CI = any(os.getenv(var) == "true" for var in ["CIRCLECI", "GITHUB_ACTIONS"])
 IN_RE_WORKER = os.environ.get("INSIDE_RE_WORKER") is not None
 IN_FBCODE = os.environ.get("IN_FBCODE_TORCHVISION") == "1"
+CVCUDA_AVAILABLE = _is_cvcuda_available()
 CUDA_NOT_AVAILABLE_MSG = "CUDA device not available"
 MPS_NOT_AVAILABLE_MSG = "MPS device not available"
 OSS_CI_GPU_NO_CUDA_MSG = "We're in an OSS GPU machine, and this test doesn't need cuda."
@@ -307,11 +303,10 @@ def __init__(
         if isinstance(expected, PIL.Image.Image):
             expected = to_image(expected)
 
-        # attempt to convert CV-CUDA tensors to torch tensors
-        if is_cvcuda_tensor(actual):
+        # handle check for CV-CUDA Tensors
+        if CVCUDA_AVAILABLE and isinstance(actual, _import_cvcuda().Tensor):
+            # Use the PIL compatible tensor, so we can always compare with PIL.Image.Image
             actual = cvcuda_to_pil_compatible_tensor(actual)
-        if is_cvcuda_tensor(expected):
-            expected = cvcuda_to_pil_compatible_tensor(expected)
 
         super().__init__(actual, expected, **other_parameters)
         self.mae = mae

test/test_transforms_v2.py

@@ -21,6 +21,7 @@
 import torchvision.transforms.v2 as transforms
 
 from common_utils import (
+    assert_close,
     assert_equal,
     cache,
     cpu_and_cuda,
@@ -42,7 +43,6 @@
 )
 
 from torch import nn
-from torch.testing import assert_close
 from torch.utils._pytree import tree_flatten, tree_map
 from torch.utils.data import DataLoader, default_collate
 from torchvision import tv_tensors
@@ -3499,11 +3499,8 @@ def test_functional_image_correctness(self, kwargs, make_input):
 
         actual = F.crop(image, **kwargs)
 
-        if make_input == make_image_cvcuda:
-            actual = F.cvcuda_to_tensor(actual).to(device="cpu")
-            actual = actual.squeeze(0)
-            image = F.cvcuda_to_tensor(image).to(device="cpu")
-            image = image.squeeze(0)
+        if make_input is make_image_cvcuda:
+            image = cvcuda_to_pil_compatible_tensor(image)
 
         expected = F.to_image(F.crop(F.to_pil_image(image), **kwargs))
 
@@ -3624,15 +3621,15 @@ def test_transform_image_correctness(self, param, value, seed, make_input):
 
             torch.manual_seed(seed)
 
-            if make_input == make_image_cvcuda:
+            if make_input is make_image_cvcuda:
                 image = cvcuda_to_pil_compatible_tensor(image)
 
             expected = F.to_image(transform(F.to_pil_image(image)))
 
         if make_input == make_image_cvcuda and will_pad:
             # when padding is applied, CV-CUDA will always fill with zeros
             # cannot use assert_equal since it will fail unless random is all zeros
-            torch.testing.assert_close(actual, expected, rtol=0, atol=get_max_value(image.dtype))
+            assert_close(actual, expected, rtol=0, atol=get_max_value(image.dtype))
         else:
             assert_equal(actual, expected)
 
@@ -4458,6 +4455,9 @@ def test_kernel(self, kernel, make_input):
             make_segmentation_mask,
             make_video,
             make_keypoints,
+            pytest.param(
+                make_image_cvcuda, marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA")
+            ),
         ],
     )
     def test_functional(self, make_input):
@@ -4474,9 +4474,16 @@ def test_functional(self, make_input):
             (F.resized_crop_mask, tv_tensors.Mask),
             (F.resized_crop_video, tv_tensors.Video),
             (F.resized_crop_keypoints, tv_tensors.KeyPoints),
+            pytest.param(
+                F.resized_crop_image,
+                "cvcuda.Tensor",
+                marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA"),
+            ),
         ],
     )
     def test_functional_signature(self, kernel, input_type):
+        if input_type == "cvcuda.Tensor":
+            input_type = _import_cvcuda().Tensor
         check_functional_kernel_signature_match(F.resized_crop, kernel=kernel, input_type=input_type)
 
     @param_value_parametrization(
@@ -4493,6 +4500,9 @@ def test_functional_signature(self, kernel, input_type):
             make_segmentation_mask,
             make_video,
             make_keypoints,
+            pytest.param(
+                make_image_cvcuda, marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA")
+            ),
         ],
     )
     def test_transform(self, param, value, make_input):
@@ -4504,20 +4514,37 @@ def test_transform(self, param, value, make_input):
 
     # `InterpolationMode.NEAREST` is modeled after the buggy `INTER_NEAREST` interpolation of CV2.
     # The PIL equivalent of `InterpolationMode.NEAREST` is `InterpolationMode.NEAREST_EXACT`
+    @pytest.mark.parametrize(
+        "make_input",
+        [
+            make_image,
+            pytest.param(
+                make_image_cvcuda, marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA")
+            ),
+        ],
+    )
     @pytest.mark.parametrize("interpolation", set(INTERPOLATION_MODES) - {transforms.InterpolationMode.NEAREST})
-    def test_functional_image_correctness(self, interpolation):
-        image = make_image(self.INPUT_SIZE, dtype=torch.uint8)
+    def test_functional_image_correctness(self, make_input, interpolation):
+        image = make_input(self.INPUT_SIZE, dtype=torch.uint8)
 
         actual = F.resized_crop(
             image, **self.CROP_KWARGS, size=self.OUTPUT_SIZE, interpolation=interpolation, antialias=True
         )
+
+        if make_input is make_image_cvcuda:
+            image = cvcuda_to_pil_compatible_tensor(image)
+
         expected = F.to_image(
             F.resized_crop(
                 F.to_pil_image(image), **self.CROP_KWARGS, size=self.OUTPUT_SIZE, interpolation=interpolation
             )
         )
 
-        torch.testing.assert_close(actual, expected, atol=1, rtol=0)
+        atol = 1
+        if make_input is make_image_cvcuda and interpolation == transforms.InterpolationMode.BICUBIC:
+            # CV-CUDA BICUBIC differs from PIL ground truth BICUBIC
+            atol = 10
+        assert_close(actual, expected, atol=atol, rtol=0)
 
     def _reference_resized_crop_bounding_boxes(self, bounding_boxes, *, top, left, height, width, size):
         new_height, new_width = size
@@ -4992,7 +5019,7 @@ def test_image_correctness(self, output_size, make_input, fn):
 
         actual = fn(image, output_size)
 
-        if make_input == make_image_cvcuda:
+        if make_input is make_image_cvcuda:
             image = cvcuda_to_pil_compatible_tensor(image)
 
         expected = F.to_image(F.center_crop(F.to_pil_image(image), output_size=output_size))

torchvision/transforms/v2/_geometry.py

@@ -46,6 +46,8 @@ class RandomHorizontalFlip(_RandomApplyTransform):
 
     _v1_transform_cls = _transforms.RandomHorizontalFlip
 
+    _transformed_types = _RandomApplyTransform._transformed_types + (is_cvcuda_tensor,)
+
     def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
         return self._call_kernel(F.horizontal_flip, inpt)
 
@@ -64,6 +66,8 @@ class RandomVerticalFlip(_RandomApplyTransform):
 
     _v1_transform_cls = _transforms.RandomVerticalFlip
 
+    _transformed_types = _RandomApplyTransform._transformed_types + (is_cvcuda_tensor,)
+
     def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
         return self._call_kernel(F.vertical_flip, inpt)
 
@@ -247,6 +251,8 @@ class RandomResizedCrop(Transform):
 
     _v1_transform_cls = _transforms.RandomResizedCrop
 
+    _transformed_types = Transform._transformed_types + (is_cvcuda_tensor,)
+
     def __init__(
         self,
         size: Union[int, Sequence[int]],

torchvision/transforms/v2/functional/_geometry.py

@@ -605,6 +605,32 @@ def resize_video(
     return resize_image(video, size=size, interpolation=interpolation, max_size=max_size, antialias=antialias)
 
 
+def _resize_cvcuda(
+    image: "cvcuda.Tensor",
+    size: Optional[list[int]],
+    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
+    max_size: Optional[int] = None,
+    antialias: Optional[bool] = True,
+) -> "cvcuda.Tensor":
+    # placeholder func for now, will be handled in PR for resize alone
+    # since placeholder convert to from torch tensor and use resize_image
+    from ._type_conversion import cvcuda_to_tensor, to_cvcuda_tensor
+
+    return to_cvcuda_tensor(
+        resize_image(
+            cvcuda_to_tensor(image),
+            size=size,
+            interpolation=interpolation,
+            max_size=max_size,
+            antialias=antialias,
+        )
+    )
+
+
+if CVCUDA_AVAILABLE:
+    _register_kernel_internal(resize, _import_cvcuda().Tensor)(_resize_cvcuda)
+
+
 def affine(
     inpt: torch.Tensor,
     angle: Union[int, float],
@@ -2946,6 +2972,24 @@ def resized_crop_video(
     )
 
 
+def _resized_crop_cvcuda(
+    image: "cvcuda.Tensor",
+    top: int,
+    left: int,
+    height: int,
+    width: int,
+    size: list[int],
+    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
+    antialias: Optional[bool] = True,
+) -> "cvcuda.Tensor":
+    image = _crop_cvcuda(image, top, left, height, width)
+    return _resize_cvcuda(image, size, interpolation=interpolation, antialias=antialias)
+
+
+if CVCUDA_AVAILABLE:
+    _register_kernel_internal(resized_crop, _import_cvcuda().Tensor)(_resized_crop_cvcuda)
+
+
 def five_crop(
     inpt: torch.Tensor, size: list[int]
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
@@ -3024,15 +3068,15 @@ def _five_crop_cvcuda(
     size: list[int],
 ) -> tuple["cvcuda.Tensor", "cvcuda.Tensor", "cvcuda.Tensor", "cvcuda.Tensor", "cvcuda.Tensor"]:
     crop_height, crop_width = _parse_five_crop_size(size)
-    image_height, image_width = image.shape[-2:]
+    image_height, image_width = image.shape[1], image.shape[2]
 
     if crop_width > image_width or crop_height > image_height:
         raise ValueError(f"Requested crop size {size} is bigger than input size {(image_height, image_width)}")
 
     tl = _crop_cvcuda(image, 0, 0, crop_height, crop_width)
-    tr = _crop_cvcuda(image, 0, image_width - crop_height, crop_width, crop_height)
-    bl = _crop_cvcuda(image, image_height - crop_height, 0, crop_width, crop_height)
-    br = _crop_cvcuda(image, image_height - crop_height, image_width - crop_width, crop_width, crop_height)
+    tr = _crop_cvcuda(image, 0, image_width - crop_width, crop_height, crop_width)
+    bl = _crop_cvcuda(image, image_height - crop_height, 0, crop_height, crop_width)
+    br = _crop_cvcuda(image, image_height - crop_height, image_width - crop_width, crop_height, crop_width)
     center = _center_crop_cvcuda(image, [crop_height, crop_width])
 
     return tl, tr, bl, br, center