CVCUDA backend design (pytorch#9259)

Summary:
Summary
-------

This PR provides the first building blocks for CV-CUDA integration in torchvision. We add the functionals `to_cvcuda_tensor` and `cvcuda_to_tensor` to transform from `torch.Tensor` to `cvcuda.Tensor` and back. We also implement the corresponding class transforms `ToCVCUDATensor` and `CVCUDAToTensor`.

**Key features:**

*   **3-channel RGB support only**: Simplified API focusing on the most common use case (RGB images)
*   **Supported data types**: `torch.uint8` (RGB8 format) and `torch.float32` (RGBf32 format)
*   **Lossless round-trip conversions**: Exact data preservation when converting PyTorch ↔ CV-CUDA
*   **Informative error messages**: Helpful installation instructions when CV-CUDA is not available
*   **Batch-aware**: Handles both unbatched (CHW) and batched (NCHW) tensors

Users must explicitly opt-in for these transforms, which require CV-CUDA to be installed.

How to use
----------

```python
from PIL import Image
import torchvision.transforms.v2.functional as F 

# Load and convert image to PyTorch tensor
orig_img = Image.open("leaning_tower.jpg")
img_tensor = F.pil_to_tensor(orig_img)  

# Convert to CV-CUDA tensor (must be 3-channel RGB on CUDA)
cvcuda_tensor = F.to_cvcuda_tensor(img_tensor.cuda())

# Convert back to PyTorch tensor
img_tensor = F.cvcuda_to_tensor(cvcuda_tensor)
```

> [!NOTE]
> 
> *   NVCV tensors are automatically converted to NHWC layout, contrary to torchvision's NCHW default
> *   Only 3-channel RGB images and 1-channel grayscale are supported for now
> *   Input tensors will be uploaded to CUDA device when converting to CV-CUDA tensors
> *   CV-CUDA must be installed: `pip install cvcuda-cu12` (CUDA 12) or `pip install cvcuda-cu11` (CUDA 11)

Run unit tests
--------------
## Run unit tests

```bash
pytest test/test_transforms_v2.py -k "cvcuda"
...
35 passed, 4 skipped, 9774 deselected in 1.12s
```


Test Plan:
```python
from torchvision import _is_cvcuda_available

_is_cvcuda_available()
```

## Run tests

```bash
buck test fbcode//mode/opt fbcode//pytorch/vision/test:torchvision_transforms_v2
...
Tests finished: Pass 38. Fail 0. Fatal 0. Skip 0. Build failure 0
```

Differential Revision: D85862362

Pulled By: AntoineSimoulin

Author: AntoineSimoulin

test/common_utils.py

@@ -20,7 +20,7 @@
 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_image, to_pil_image
+from torchvision.transforms.v2.functional import to_cvcuda_tensor, to_image, to_pil_image
 from torchvision.utils import _Image_fromarray
 
 
@@ -400,6 +400,10 @@ def make_image_pil(*args, **kwargs):
     return to_pil_image(make_image(*args, **kwargs))
 
 
+def make_image_cvcuda(*args, **kwargs):
+    return to_cvcuda_tensor(make_image(*args, **kwargs))
+
+
 def make_keypoints(canvas_size=DEFAULT_SIZE, *, num_points=4, dtype=None, device="cpu"):
     y = torch.randint(0, canvas_size[0], size=(num_points, 1), dtype=dtype, device=device)
     x = torch.randint(0, canvas_size[1], size=(num_points, 1), dtype=dtype, device=device)

test/test_transforms_v2.py

@@ -29,6 +29,7 @@
     make_bounding_boxes,
     make_detection_masks,
     make_image,
+    make_image_cvcuda,
     make_image_pil,
     make_image_tensor,
     make_keypoints,
@@ -51,8 +52,16 @@
 from torchvision.transforms.v2 import functional as F
 from torchvision.transforms.v2._utils import check_type, is_pure_tensor
 from torchvision.transforms.v2.functional._geometry import _get_perspective_coeffs, _parallelogram_to_bounding_boxes
+from torchvision.transforms.v2.functional._type_conversion import _import_cvcuda_modules
 from torchvision.transforms.v2.functional._utils import _get_kernel, _register_kernel_internal
 
+try:
+    _import_cvcuda_modules()
+    CVCUDA_AVAILABLE = True
+except ImportError:
+    CVCUDA_AVAILABLE = False
+CUDA_AVAILABLE = torch.cuda.is_available()
+
 
 # turns all warnings into errors for this module
 pytestmark = [pytest.mark.filterwarnings("error")]
@@ -6733,6 +6742,184 @@ def test_functional_error(self):
             F.pil_to_tensor(object())
 
 
+@pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA")
+@pytest.mark.skipif(not CUDA_AVAILABLE, reason="test requires CUDA")
+class TestToCVCUDATensor:
+    """Tests for to_cvcuda_tensor function following patterns from TestToPil"""
+
+    def test_1_channel_uint8_tensor_to_cvcuda_tensor(self):
+        # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
+        img_data = torch.randint(0, 256, (1, 4, 4), dtype=torch.uint8)
+        img_data = img_data.cuda()
+        cvcuda_img = F.to_cvcuda_tensor(img_data)
+        # Check that the conversion succeeded and format is correct
+        assert cvcuda_img is not None
+
+    def test_1_channel_int16_tensor_to_cvcuda_tensor(self):
+        # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
+        img_data = torch.randint(0, 256, (1, 4, 4), dtype=torch.int16)
+        img_data = img_data.cuda()
+        cvcuda_img = F.to_cvcuda_tensor(img_data)
+        assert cvcuda_img is not None
+
+    def test_1_channel_int32_tensor_to_cvcuda_tensor(self):
+        # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
+        img_data = torch.randint(0, 256, (1, 4, 4), dtype=torch.int32)
+        img_data = img_data.cuda()
+        cvcuda_img = F.to_cvcuda_tensor(img_data)
+        assert cvcuda_img is not None
+
+    def test_1_channel_float32_tensor_to_cvcuda_tensor(self):
+        # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
+        img_data = torch.rand(1, 4, 4)
+        img_data = img_data.cuda()
+        cvcuda_img = F.to_cvcuda_tensor(img_data)
+        assert cvcuda_img is not None
+
+    def test_3_channel_uint8_tensor_to_cvcuda_tensor(self):
+        # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
+        img_data = torch.randint(0, 256, (3, 4, 4), dtype=torch.uint8)
+        img_data = img_data.cuda()
+        cvcuda_img = F.to_cvcuda_tensor(img_data)
+        assert cvcuda_img is not None
+
+    def test_3_channel_float32_tensor_to_cvcuda_tensor(self):
+        # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
+        img_data = torch.rand(3, 4, 4)
+        img_data = img_data.cuda()
+        cvcuda_img = F.to_cvcuda_tensor(img_data)
+        assert cvcuda_img is not None
+
+    def test_unsupported_num_channels(self):
+        # Test 2-channel image (not supported)
+        # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
+        img_data = torch.rand(2, 5, 5)
+        img_data = img_data.cuda()
+        with pytest.raises(ValueError, match="Only 1 and 3 channel images are supported"):
+            F.to_cvcuda_tensor(img_data)
+
+        # Test 4-channel image (not supported)
+        img_data = torch.randint(0, 256, (4, 5, 5), dtype=torch.uint8)
+        img_data = img_data.cuda()
+        with pytest.raises(ValueError, match="Only 1 and 3 channel images are supported"):
+            F.to_cvcuda_tensor(img_data)
+
+    def test_invalid_input_type(self):
+        with pytest.raises(TypeError, match=r"pic should be `torch.Tensor`"):
+            F.to_cvcuda_tensor("invalid_input")
+
+    def test_invalid_dimensions(self):
+        # Test 1D array (too few dimensions)
+        # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
+        with pytest.raises(ValueError, match=r"pic should be 3 or 4 dimensional"):
+            img_data = torch.randint(0, 256, (4,), dtype=torch.uint8)
+            img_data = img_data.cuda()
+            F.to_cvcuda_tensor(img_data)
+
+        # Test 2D array (no longer supported)
+        with pytest.raises(ValueError, match=r"pic should be 3 or 4 dimensional"):
+            img_data = torch.randint(0, 256, (4, 4), dtype=torch.uint8)
+            img_data = img_data.cuda()
+            F.to_cvcuda_tensor(img_data)
+
+        # Test 5D array (too many dimensions)
+        with pytest.raises(ValueError, match=r"pic should be 3 or 4 dimensional"):
+            img_data = torch.randint(0, 256, (1, 1, 3, 4, 4), dtype=torch.uint8)
+            img_data = img_data.cuda()
+            F.to_cvcuda_tensor(img_data)
+
+    def test_float64_tensor_to_cvcuda_tensor(self):
+        # Test single channel float64 (F64 format is supported)
+        # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
+        img_data = torch.rand(1, 4, 4, dtype=torch.float64)
+        img_data = img_data.cuda()
+        cvcuda_img = F.to_cvcuda_tensor(img_data)
+        assert cvcuda_img is not None
+
+    def test_float64_rgb_not_supported(self):
+        # Test 3-channel float64 is NOT supported (no RGBf64 format in CV-CUDA)
+        # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
+        img_data = torch.rand(3, 4, 4, dtype=torch.float64)
+        img_data = img_data.cuda()
+        with pytest.raises(TypeError, match=r"Unsupported dtype"):
+            F.to_cvcuda_tensor(img_data)
+
+    @pytest.mark.parametrize("num_channels", [1, 3])
+    @pytest.mark.parametrize("dtype", [torch.uint8, torch.float32, torch.float64])
+    def test_round_trip(self, num_channels, dtype):
+        # Skip float64 for 3-channel (not supported by CV-CUDA)
+        if num_channels == 3 and dtype == torch.float64:
+            pytest.skip("float64 is not supported for 3-channel RGB images")
+
+        # Setup: Create a tensor in CHW format (PyTorch standard)
+        # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
+        if dtype == torch.uint8:
+            original_tensor = torch.randint(0, 256, (num_channels, 4, 4), dtype=dtype)
+        else:
+            original_tensor = torch.rand(num_channels, 4, 4, dtype=dtype)
+        original_tensor = original_tensor.cuda()
+
+        # Execute: Convert to CV-CUDA and back to tensor
+        # CHW -> (to_cvcuda_tensor) -> CV-CUDA NHWC -> (cvcuda_to_tensor) -> NCHW
+        cvcuda_tensor = F.to_cvcuda_tensor(original_tensor)
+        result_tensor = F.cvcuda_to_tensor(cvcuda_tensor)
+
+        # Remove batch dimension that was added during conversion since original was unbatched
+        result_tensor = result_tensor.squeeze(0)
+
+        # Assert: The round-trip conversion preserves the original tensor exactly
+        torch.testing.assert_close(result_tensor, original_tensor, rtol=0, atol=0)
+
+    @pytest.mark.parametrize("num_channels", [1, 3])
+    @pytest.mark.parametrize("dtype", [torch.uint8, torch.float32, torch.float64])
+    @pytest.mark.parametrize("batch_size", [1, 2, 4])
+    def test_round_trip_batched(self, num_channels, dtype, batch_size):
+        # Skip float64 for 3-channel (not supported by CV-CUDA)
+        if num_channels == 3 and dtype == torch.float64:
+            pytest.skip("float64 is not supported for 3-channel RGB images")
+
+        # Setup: Create a batched tensor in NCHW format
+        # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
+        if dtype == torch.uint8:
+            original_tensor = torch.randint(0, 256, (batch_size, num_channels, 4, 4), dtype=dtype)
+        else:
+            original_tensor = torch.rand(batch_size, num_channels, 4, 4, dtype=dtype)
+        original_tensor = original_tensor.cuda()
+
+        # Execute: Convert to CV-CUDA and back to tensor
+        # NCHW -> (to_cvcuda_tensor) -> CV-CUDA NHWC -> (cvcuda_to_tensor) -> NCHW
+        cvcuda_tensor = F.to_cvcuda_tensor(original_tensor)
+        result_tensor = F.cvcuda_to_tensor(cvcuda_tensor)
+
+        # Assert: The round-trip conversion preserves the original batched tensor exactly
+        torch.testing.assert_close(result_tensor, original_tensor, rtol=0, atol=0)
+        # Also verify batch size is preserved
+        assert result_tensor.shape[0] == batch_size
+
+
+@pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA")
+@pytest.mark.skipif(not CUDA_AVAILABLE, reason="test requires CUDA")
+class TestCVDUDAToTensor:
+    @pytest.mark.parametrize("color_space", ["RGB", "GRAY"])
+    @pytest.mark.parametrize(
+        "fn",
+        [F.cvcuda_to_tensor, transform_cls_to_functional(transforms.CVCUDAToTensor)],
+    )
+    def test_functional_and_transform(self, color_space, fn):
+        input = make_image_cvcuda(color_space=color_space)
+
+        output = fn(input)
+
+        assert isinstance(output, torch.Tensor)
+        # Convert input to tensor to compare sizes
+        input_tensor = F.cvcuda_to_tensor(input)
+        assert F.get_size(output) == F.get_size(input_tensor)
+
+    def test_functional_error(self):
+        with pytest.raises(TypeError, match="cvcuda_img should be `cvcuda.Tensor`"):
+            F.cvcuda_to_tensor(object())
+
+
 class TestLambda:
     @pytest.mark.parametrize("input", [object(), torch.empty(()), np.empty(()), "string", 1, 0.0])
     @pytest.mark.parametrize("types", [(), (torch.Tensor, np.ndarray)])

torchvision/transforms/v2/__init__.py

@@ -55,7 +55,7 @@
     ToDtype,
 )
 from ._temporal import UniformTemporalSubsample
-from ._type_conversion import PILToTensor, ToImage, ToPILImage, ToPureTensor
+from ._type_conversion import CVCUDAToTensor, PILToTensor, ToCVCUDATensor, ToImage, ToPILImage, ToPureTensor
 from ._utils import check_type, get_bounding_boxes, get_keypoints, has_all, has_any, query_chw, query_size
 
 from ._deprecated import ToTensor  # usort: skip

torchvision/transforms/v2/_type_conversion.py

@@ -6,8 +6,8 @@
 
 from torchvision import tv_tensors
 from torchvision.transforms.v2 import functional as F, Transform
-
 from torchvision.transforms.v2._utils import is_pure_tensor
+from torchvision.utils import _log_api_usage_once
 
 
 class PILToTensor(Transform):
@@ -90,3 +90,71 @@ class ToPureTensor(Transform):
 
     def transform(self, inpt: Any, params: dict[str, Any]) -> torch.Tensor:
         return inpt.as_subclass(torch.Tensor)
+
+
+class ToCVCUDATensor:
+    """Convert a torch.Tensor to cvcuda.Tensor
+
+    This transform does not support torchscript.
+
+    Converts a torch.*Tensor of shape C x H x W to a cvcuda.Tensor.
+    Only 1-channel and 3-channel images are supported.
+    """
+
+    def __init__(self):
+        _log_api_usage_once(self)
+
+    def __call__(self, pic):
+        """
+        Args:
+            pic (torch.Tensor): Image to be converted to cvcuda.Tensor.
+
+        Returns:
+            cvcuda.Tensor: Image converted to cvcuda.Tensor.
+
+        """
+        return F.to_cvcuda_tensor(pic)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}()"
+
+
+class CVCUDAToTensor:
+    """Convert a `cvcuda.Tensor` to a `torch.Tensor` of the same type - this does not scale values.
+
+    This transform does not support torchscript.
+
+    Converts a `cvcuda.Tensor` to a `torch.Tensor`. Supports both batched and unbatched inputs:
+    - Unbatched: (H, W, C) or (H, W) → (C, H, W) or (1, H, W)
+    - Batched: (N, H, W, C) or (N, H, W) → (N, C, H, W) or (N, 1, H, W)
+
+    The conversion happens directly on GPU when the `cvcuda.Tensor` is stored on GPU,
+    avoiding unnecessary data transfers.
+
+    Example:
+        >>> import cvcuda
+        >>> import torchvision.transforms.v2 as T
+        >>> # Create a CV-CUDA tensor (320x240 RGB)
+        >>> # Note: In CV-CUDA 0.16.0+, Image/Tensor creation uses cvcuda module
+        >>> img_tensor = torch.randint(0, 255, (1, 240, 320, 3), dtype=torch.uint8, device="cuda")
+        >>> cvcuda_tensor = cvcuda.as_tensor(img_tensor, cvcuda.TensorLayout.NHWC)
+        >>> tensor = T.CVCUDAToTensor()(cvcuda_tensor)
+        >>> print(tensor.shape)
+        torch.Size([1, 3, 240, 320])
+    """
+
+    def __init__(self) -> None:
+        _log_api_usage_once(self)
+
+    def __call__(self, pic):
+        """
+        Args:
+            pic (cvcuda.Tensor): CV-CUDA Tensor to be converted to tensor.
+
+        Returns:
+            Tensor: Converted image in CHW format.
+        """
+        return F.cvcuda_to_tensor(pic)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}()"

torchvision/transforms/v2/functional/__init__.py

@@ -162,6 +162,6 @@
     to_dtype_video,
 )
 from ._temporal import uniform_temporal_subsample, uniform_temporal_subsample_video
-from ._type_conversion import pil_to_tensor, to_image, to_pil_image
+from ._type_conversion import cvcuda_to_tensor, pil_to_tensor, to_cvcuda_tensor, to_image, to_pil_image
 
 from ._deprecated import get_image_size, to_tensor  # usort: skip