add helper function _import_cvcuda_modules for better logging

Author: AntoineSimoulin

torchvision/transforms/v2/functional/_type_conversion.py

@@ -11,6 +11,29 @@
     import nvcv  # type: ignore[import-not-found]
 
 
+def _import_cvcuda_modules():
+    """Import CV-CUDA modules with informative error message if not installed.
+
+    Returns:
+        tuple: (cvcuda, nvcv) modules.
+
+    Raises:
+        RuntimeError: If CV-CUDA is not installed.
+    """
+    try:
+        import cvcuda  # type: ignore[import-not-found]
+        import nvcv  # type: ignore[import-not-found]
+
+        return cvcuda, nvcv
+    except ImportError as e:
+        raise RuntimeError(
+            "CV-CUDA is required but not installed. "
+            "Please install it following the instructions at "
+            "https://github.com/CVCUDA/CV-CUDA or via pip: "
+            "`pip install cvcuda-cu12` (for CUDA 12) or `pip install cvcuda-cu11` (for CUDA 11)."
+        ) from e
+
+
 @torch.jit.unused
 def to_image(inpt: Union[torch.Tensor, PIL.Image.Image, np.ndarray]) -> tv_tensors.Image:
     """See :class:`~torchvision.transforms.v2.ToImage` for details."""
@@ -45,7 +68,7 @@ def _infer_nvcv_format(img_tensor: torch.Tensor):
         TypeError: If dtype is not supported.
         ValueError: If number of channels is not 3.
     """
-    import nvcv  # type: ignore[import-not-found]
+    _, nvcv = _import_cvcuda_modules()
 
     num_channels = img_tensor.shape[2]
     dtype = img_tensor.dtype
@@ -77,8 +100,7 @@ def to_nvcv_tensor(pic) -> "nvcv.Tensor":
     Returns:
         nvcv.Tensor: Image converted to nvcv.Tensor with NHWC layout.
     """
-    import cvcuda  # type: ignore[import-not-found]
-    import nvcv  # type: ignore[import-not-found]
+    cvcuda, nvcv = _import_cvcuda_modules()
 
     if not torch.jit.is_scripting() and not torch.jit.is_tracing():
         _log_api_usage_once(to_nvcv_tensor)
@@ -119,7 +141,7 @@ def nvcv_to_tensor(nvcv_img: "nvcv.Tensor") -> torch.Tensor:
     Returns:
         torch.Tensor: Converted image in NCHW format (batched).
     """
-    import nvcv  # type: ignore[import-not-found]
+    _, nvcv = _import_cvcuda_modules()
 
     if not torch.jit.is_scripting() and not torch.jit.is_tracing():
         _log_api_usage_once(nvcv_to_tensor)