begin work on finalizing the crop PR to include five and ten crop, adhere to new PR reviews for flip

Author: justincdavis

test/common_utils.py

@@ -20,7 +20,13 @@
 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 to_cvcuda_tensor, to_image, to_pil_image
+from torchvision.transforms.v2.functional import (
+    cvcuda_to_tensor,
+    is_cvcuda_tensor,
+    to_cvcuda_tensor,
+    to_image,
+    to_pil_image,
+)
 from torchvision.utils import _Image_fromarray
 
 
@@ -275,6 +281,17 @@ def combinations_grid(**kwargs):
     return [dict(zip(kwargs.keys(), values)) for values in itertools.product(*kwargs.values())]
 
 
+def cvcuda_to_pil_compatible_tensor(tensor: "cvcuda.Tensor") -> torch.Tensor:
+    tensor = cvcuda_to_tensor(tensor)
+    if tensor.ndim != 4:
+        raise ValueError(f"CV-CUDA Tensor should be 4 dimensional. Got {tensor.ndim} dimensions.")
+    if tensor.shape[0] != 1:
+        raise ValueError(
+            f"CV-CUDA Tensor should have batch dimension 1 for comparison with PIL.Image.Image. Got {tensor.shape[0]}."
+        )
+    return tensor.squeeze(0).cpu()
+
+
 class ImagePair(TensorLikePair):
     def __init__(
         self,
@@ -284,8 +301,17 @@ def __init__(
         mae=False,
         **other_parameters,
     ):
-        if all(isinstance(input, PIL.Image.Image) for input in [actual, expected]):
-            actual, expected = (to_image(input) for input in [actual, expected])
+        # Convert PIL images to tv_tensors.Image (regardless of what the other is)
+        if isinstance(actual, PIL.Image.Image):
+            actual = to_image(actual)
+        if isinstance(expected, PIL.Image.Image):
+            expected = to_image(expected)
+
+        # attempt to convert CV-CUDA tensors to torch tensors
+        if is_cvcuda_tensor(actual):
+            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
@@ -401,7 +427,6 @@ def make_image_pil(*args, **kwargs):
 
 
 def make_image_cvcuda(*args, batch_dims=(1,), **kwargs):
-    # explicitly default batch_dims to (1,) since to_cvcuda_tensor requires a batch dimension (ndims == 4)
     return to_cvcuda_tensor(make_image(*args, batch_dims=batch_dims, **kwargs))
 
 

test/test_transforms_v2.py

@@ -24,6 +24,7 @@
     assert_equal,
     cache,
     cpu_and_cuda,
+    cvcuda_to_pil_compatible_tensor,
     freeze_rng_state,
     ignore_jit_no_profile_information_warning,
     make_bounding_boxes,
@@ -3624,10 +3625,7 @@ def test_transform_image_correctness(self, param, value, seed, make_input):
             torch.manual_seed(seed)
 
             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)
+                image = cvcuda_to_pil_compatible_tensor(image)
 
             expected = F.to_image(transform(F.to_pil_image(image)))
 
@@ -4995,10 +4993,7 @@ def test_image_correctness(self, output_size, make_input, fn):
         actual = fn(image, output_size)
 
         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)
+            image = cvcuda_to_pil_compatible_tensor(image)
 
         expected = F.to_image(F.center_crop(F.to_pil_image(image), output_size=output_size))
 
@@ -6274,7 +6269,15 @@ def wrapper(*args, **kwargs):
 
     @pytest.mark.parametrize(
         "make_input",
-        [make_image_tensor, make_image_pil, make_image, make_video],
+        [
+            make_image_tensor,
+            make_image_pil,
+            make_image,
+            make_video,
+            pytest.param(
+                make_image_cvcuda, marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA")
+            ),
+        ],
     )
     @pytest.mark.parametrize("functional", [F.five_crop, F.ten_crop])
     def test_functional(self, make_input, functional):
@@ -6292,13 +6295,27 @@ def test_functional(self, make_input, functional):
             (F.five_crop, F._geometry._five_crop_image_pil, PIL.Image.Image),
             (F.five_crop, F.five_crop_image, tv_tensors.Image),
             (F.five_crop, F.five_crop_video, tv_tensors.Video),
+            pytest.param(
+                F.five_crop,
+                F._geometry._five_crop_cvcuda,
+                "cvcuda.Tensor",
+                marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA"),
+            ),
             (F.ten_crop, F.ten_crop_image, torch.Tensor),
             (F.ten_crop, F._geometry._ten_crop_image_pil, PIL.Image.Image),
             (F.ten_crop, F.ten_crop_image, tv_tensors.Image),
             (F.ten_crop, F.ten_crop_video, tv_tensors.Video),
+            pytest.param(
+                F.ten_crop,
+                F._geometry._ten_crop_cvcuda,
+                "cvcuda.Tensor",
+                marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA"),
+            ),
         ],
     )
     def test_functional_signature(self, functional, kernel, input_type):
+        if input_type == "cvcuda.Tensor":
+            input_type = _import_cvcuda().Tensor
         check_functional_kernel_signature_match(functional, kernel=kernel, input_type=input_type)
 
     class _TransformWrapper(nn.Module):
@@ -6320,7 +6337,15 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
 
     @pytest.mark.parametrize(
         "make_input",
-        [make_image_tensor, make_image_pil, make_image, make_video],
+        [
+            make_image_tensor,
+            make_image_pil,
+            make_image,
+            make_video,
+            pytest.param(
+                make_image_cvcuda, marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA")
+            ),
+        ],
     )
     @pytest.mark.parametrize("transform_cls", [transforms.FiveCrop, transforms.TenCrop])
     def test_transform(self, make_input, transform_cls):
@@ -6338,29 +6363,55 @@ def test_transform_error(self, make_input, transform_cls):
         with pytest.raises(TypeError, match="not supported"):
             transform(make_input(self.INPUT_SIZE))
 
+    @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("fn", [F.five_crop, transform_cls_to_functional(transforms.FiveCrop)])
-    def test_correctness_image_five_crop(self, fn):
-        image = make_image(self.INPUT_SIZE, dtype=torch.uint8, device="cpu")
+    def test_correctness_image_five_crop(self, make_input, fn):
+        image = make_input(self.INPUT_SIZE, dtype=torch.uint8, device="cpu")
 
         actual = fn(image, size=self.OUTPUT_SIZE)
+
+        if make_input is make_image_cvcuda:
+            image = cvcuda_to_pil_compatible_tensor(image)
+
         expected = F.five_crop(F.to_pil_image(image), size=self.OUTPUT_SIZE)
 
         assert isinstance(actual, tuple)
         assert_equal(actual, [F.to_image(e) for e in expected])
 
+    @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("fn_or_class", [F.ten_crop, transforms.TenCrop])
     @pytest.mark.parametrize("vertical_flip", [False, True])
-    def test_correctness_image_ten_crop(self, fn_or_class, vertical_flip):
+    def test_correctness_image_ten_crop(self, make_input, fn_or_class, vertical_flip):
         if fn_or_class is transforms.TenCrop:
             fn = transform_cls_to_functional(fn_or_class, size=self.OUTPUT_SIZE, vertical_flip=vertical_flip)
             kwargs = dict()
         else:
             fn = fn_or_class
             kwargs = dict(size=self.OUTPUT_SIZE, vertical_flip=vertical_flip)
 
-        image = make_image(self.INPUT_SIZE, dtype=torch.uint8, device="cpu")
+        image = make_input(self.INPUT_SIZE, dtype=torch.uint8, device="cpu")
 
         actual = fn(image, **kwargs)
+
+        if make_input is make_image_cvcuda:
+            image = cvcuda_to_pil_compatible_tensor(image)
+
         expected = F.ten_crop(F.to_pil_image(image), size=self.OUTPUT_SIZE, vertical_flip=vertical_flip)
 
         assert isinstance(actual, tuple)

torchvision/transforms/v2/_geometry.py

@@ -26,6 +26,7 @@
     get_bounding_boxes,
     has_all,
     has_any,
+    is_cvcuda_tensor,
     is_pure_tensor,
     query_size,
 )
@@ -186,6 +187,8 @@ class CenterCrop(Transform):
 
     _v1_transform_cls = _transforms.CenterCrop
 
+    _transformed_types = Transform._transformed_types + (is_cvcuda_tensor,)
+
     def __init__(self, size: Union[int, Sequence[int]]):
         super().__init__()
         self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.")
@@ -352,6 +355,8 @@ class FiveCrop(Transform):
 
     _v1_transform_cls = _transforms.FiveCrop
 
+    _transformed_types = Transform._transformed_types + (is_cvcuda_tensor,)
+
     def __init__(self, size: Union[int, Sequence[int]]) -> None:
         super().__init__()
         self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.")
@@ -396,6 +401,8 @@ class TenCrop(Transform):
 
     _v1_transform_cls = _transforms.TenCrop
 
+    _transformed_types = Transform._transformed_types + (is_cvcuda_tensor,)
+
     def __init__(self, size: Union[int, Sequence[int]], vertical_flip: bool = False) -> None:
         super().__init__()
         self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.")
@@ -803,6 +810,8 @@ class RandomCrop(Transform):
 
     _v1_transform_cls = _transforms.RandomCrop
 
+    _transformed_types = Transform._transformed_types + (is_cvcuda_tensor,)
+
     def _extract_params_for_v1_transform(self) -> dict[str, Any]:
         params = super()._extract_params_for_v1_transform()
 
@@ -1113,6 +1122,8 @@ class RandomIoUCrop(Transform):
             Default, 40.
     """
 
+    _transformed_types = Transform._transformed_types + (is_cvcuda_tensor,)
+
     def __init__(
         self,
         min_scale: float = 0.3,

torchvision/transforms/v2/_transform.py

@@ -8,10 +8,10 @@
 from torch import nn
 from torch.utils._pytree import tree_flatten, tree_unflatten
 from torchvision import tv_tensors
-from torchvision.transforms.v2._utils import check_type, has_any, is_pure_tensor
+from torchvision.transforms.v2._utils import check_type, has_any, is_cvcuda_tensor, is_pure_tensor
 from torchvision.utils import _log_api_usage_once
 
-from .functional._utils import _get_kernel, is_cvcuda_tensor
+from .functional._utils import _get_kernel
 
 
 class Transform(nn.Module):
@@ -23,7 +23,7 @@ class Transform(nn.Module):
 
     # Class attribute defining transformed types. Other types are passed-through without any transformation
     # We support both Types and callables that are able to do further checks on the type of the input.
-    _transformed_types: tuple[type | Callable[[Any], bool], ...] = (torch.Tensor, PIL.Image.Image, is_cvcuda_tensor)
+    _transformed_types: tuple[type | Callable[[Any], bool], ...] = (torch.Tensor, PIL.Image.Image)
 
     def __init__(self) -> None:
         super().__init__()

torchvision/transforms/v2/functional/_geometry.py

@@ -142,6 +142,14 @@ def horizontal_flip_video(video: torch.Tensor) -> torch.Tensor:
     return horizontal_flip_image(video)
 
 
+def _horizontal_flip_cvcuda(image: "cvcuda.Tensor") -> "cvcuda.Tensor":
+    return _import_cvcuda().flip(image, flipCode=1)
+
+
+if CVCUDA_AVAILABLE:
+    _register_kernel_internal(horizontal_flip, _import_cvcuda().Tensor)(_horizontal_flip_cvcuda)
+
+
 def vertical_flip(inpt: torch.Tensor) -> torch.Tensor:
     """See :class:`~torchvision.transforms.v2.RandomVerticalFlip` for details."""
     if torch.jit.is_scripting():
@@ -230,6 +238,14 @@ def vertical_flip_video(video: torch.Tensor) -> torch.Tensor:
     return vertical_flip_image(video)
 
 
+def _vertical_flip_cvcuda(image: "cvcuda.Tensor") -> "cvcuda.Tensor":
+    return _import_cvcuda().flip(image, flipCode=0)
+
+
+if CVCUDA_AVAILABLE:
+    _register_kernel_internal(vertical_flip, _import_cvcuda().Tensor)(_vertical_flip_cvcuda)
+
+
 # We changed the names to align them with the transforms, i.e. `RandomHorizontalFlip`. Still, `hflip` and `vflip` are
 # prevalent and well understood. Thus, we just alias them without deprecating the old names.
 hflip = horizontal_flip
@@ -3003,6 +3019,29 @@ def five_crop_video(
     return five_crop_image(video, size)
 
 
+def _five_crop_cvcuda(
+    image: "cvcuda.Tensor",
+    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:]
+
+    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)
+    center = _center_crop_cvcuda(image, [crop_height, crop_width])
+
+    return tl, tr, bl, br, center
+
+
+if CVCUDA_AVAILABLE:
+    _register_kernel_internal(five_crop, _import_cvcuda().Tensor)(_five_crop_cvcuda)
+
+
 def ten_crop(
     inpt: torch.Tensor, size: list[int], vertical_flip: bool = False
 ) -> tuple[
@@ -3098,3 +3137,35 @@ def ten_crop_video(
     torch.Tensor,
 ]:
     return ten_crop_image(video, size, vertical_flip=vertical_flip)
+
+
+def _ten_crop_cvcuda(
+    image: "cvcuda.Tensor",
+    size: list[int],
+    vertical_flip: bool = False,
+) -> tuple[
+    "cvcuda.Tensor",
+    "cvcuda.Tensor",
+    "cvcuda.Tensor",
+    "cvcuda.Tensor",
+    "cvcuda.Tensor",
+    "cvcuda.Tensor",
+    "cvcuda.Tensor",
+    "cvcuda.Tensor",
+    "cvcuda.Tensor",
+    "cvcuda.Tensor",
+]:
+    non_flipped = _five_crop_cvcuda(image, size)
+
+    if vertical_flip:
+        image = _vertical_flip_cvcuda(image)
+    else:
+        image = _horizontal_flip_cvcuda(image)
+
+    flipped = _five_crop_cvcuda(image, size)
+
+    return non_flipped + flipped
+
+
+if CVCUDA_AVAILABLE:
+    _register_kernel_internal(ten_crop, _import_cvcuda().Tensor)(_ten_crop_cvcuda)