rebase to main standards

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,
@@ -3374,14 +3373,14 @@ def test_functional(self, make_input):
             (F.elastic_video, tv_tensors.Video),
             (F.elastic_keypoints, tv_tensors.KeyPoints),
             pytest.param(
-                F._geometry._elastic_cvcuda,
-                "cvcuda.Tensor",
+                F._geometry._elastic_image_cvcuda,
+                None,
                 marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="CV-CUDA not available"),
             ),
         ],
     )
     def test_functional_signature(self, kernel, input_type):
-        if input_type == "cvcuda.Tensor":
+        if kernel is F._geometry._elastic_image_cvcuda:
             input_type = _import_cvcuda().Tensor
         check_functional_kernel_signature_match(F.elastic, kernel=kernel, input_type=input_type)
 

torchvision/transforms/v2/_geometry.py

@@ -29,7 +29,6 @@
     is_pure_tensor,
     query_size,
 )
-from .functional._utils import is_cvcuda_tensor
 
 CVCUDA_AVAILABLE = _is_cvcuda_available()
 
@@ -1046,7 +1045,8 @@ class ElasticTransform(Transform):
 
     _v1_transform_cls = _transforms.ElasticTransform
 
-    _transformed_types = Transform._transformed_types + (is_cvcuda_tensor,)
+    if CVCUDA_AVAILABLE:
+        _transformed_types = Transform._transformed_types + (_is_cvcuda_tensor,)
 
     def __init__(
         self,

torchvision/transforms/v2/functional/_augment.py

@@ -1,5 +1,4 @@
 import io
-from typing import TYPE_CHECKING
 
 import PIL.Image
 
@@ -9,15 +8,7 @@
 from torchvision.transforms.functional import pil_to_tensor, to_pil_image
 from torchvision.utils import _log_api_usage_once
 
-from ._utils import _get_kernel, _import_cvcuda, _is_cvcuda_available, _register_kernel_internal
-
-
-CVCUDA_AVAILABLE = _is_cvcuda_available()
-
-if TYPE_CHECKING:
-    import cvcuda  # type: ignore[import-not-found]
-if CVCUDA_AVAILABLE:
-    cvcuda = _import_cvcuda()  # noqa: F811
+from ._utils import _get_kernel, _register_kernel_internal
 
 
 def erase(

torchvision/transforms/v2/functional/_color.py

@@ -1,5 +1,3 @@
-from typing import TYPE_CHECKING
-
 import PIL.Image
 import torch
 from torch.nn.functional import conv2d
@@ -11,15 +9,7 @@
 
 from ._misc import _num_value_bits, to_dtype_image
 from ._type_conversion import pil_to_tensor, to_pil_image
-from ._utils import _get_kernel, _import_cvcuda, _is_cvcuda_available, _register_kernel_internal
-
-
-CVCUDA_AVAILABLE = _is_cvcuda_available()
-
-if TYPE_CHECKING:
-    import cvcuda  # type: ignore[import-not-found]
-if CVCUDA_AVAILABLE:
-    cvcuda = _import_cvcuda()  # noqa: F811
+from ._utils import _get_kernel, _register_kernel_internal
 
 
 def rgb_to_grayscale(inpt: torch.Tensor, num_output_channels: int = 1) -> torch.Tensor:

torchvision/transforms/v2/functional/_geometry.py

@@ -29,6 +29,7 @@
 
 from ._utils import (
     _FillTypeJIT,
+    _get_cvcuda_interp,
     _get_kernel,
     _import_cvcuda,
     _is_cvcuda_available,
@@ -2530,36 +2531,14 @@ def elastic_video(
     return elastic_image(video, displacement, interpolation=interpolation, fill=fill)
 
 
-if CVCUDA_AVAILABLE:
-    _cvcuda_interp = {
-        InterpolationMode.BILINEAR: cvcuda.Interp.LINEAR,
-        "bilinear": cvcuda.Interp.LINEAR,
-        "linear": cvcuda.Interp.LINEAR,
-        2: cvcuda.Interp.LINEAR,
-        InterpolationMode.BICUBIC: cvcuda.Interp.CUBIC,
-        "bicubic": cvcuda.Interp.CUBIC,
-        3: cvcuda.Interp.CUBIC,
-        InterpolationMode.NEAREST: cvcuda.Interp.NEAREST,
-        "nearest": cvcuda.Interp.NEAREST,
-        0: cvcuda.Interp.NEAREST,
-        InterpolationMode.BOX: cvcuda.Interp.BOX,
-        "box": cvcuda.Interp.BOX,
-        4: cvcuda.Interp.BOX,
-        InterpolationMode.HAMMING: cvcuda.Interp.HAMMING,
-        "hamming": cvcuda.Interp.HAMMING,
-        5: cvcuda.Interp.HAMMING,
-        InterpolationMode.LANCZOS: cvcuda.Interp.LANCZOS,
-        "lanczos": cvcuda.Interp.LANCZOS,
-        1: cvcuda.Interp.LANCZOS,
-    }
-
-
-def _elastic_cvcuda(
+def _elastic_image_cvcuda(
     image: "cvcuda.Tensor",
     displacement: torch.Tensor,
     interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
     fill: _FillTypeJIT = None,
 ) -> "cvcuda.Tensor":
+    cvcuda = _import_cvcuda()
+
     if not isinstance(displacement, torch.Tensor):
         raise TypeError("Argument displacement should be a Tensor")
 
@@ -2578,9 +2557,7 @@ def _elastic_cvcuda(
     elif num_channels == 1 and input_dtype != cvcuda.Type.F32:
         raise ValueError(f"cvcuda.remap requires float32 dtype for 1-channel images, but got {input_dtype}")
 
-    interp = _cvcuda_interp.get(interpolation, cvcuda.Interp.LINEAR)
-    if interp is None:
-        raise ValueError(f"Invalid interpolation mode: {interpolation}")
+    interp = _get_cvcuda_interp(interpolation)
 
     # Build normalized grid: identity + displacement
     # _create_identity_grid returns (1, H, W, 2) with values in [-1, 1]
@@ -2627,7 +2604,7 @@ def _elastic_cvcuda(
 
 
 if CVCUDA_AVAILABLE:
-    _elastic_cvcuda = _register_kernel_internal(elastic, cvcuda.Tensor)(_elastic_cvcuda)
+    _register_kernel_internal(elastic, _import_cvcuda().Tensor)(_elastic_image_cvcuda)
 
 
 def center_crop(inpt: torch.Tensor, output_size: list[int]) -> torch.Tensor:

torchvision/transforms/v2/functional/_misc.py

@@ -1,5 +1,5 @@
 import math
-from typing import Optional, TYPE_CHECKING
+from typing import Optional
 
 import PIL.Image
 import torch
@@ -13,14 +13,7 @@
 
 from ._meta import _convert_bounding_box_format
 
-from ._utils import _get_kernel, _import_cvcuda, _is_cvcuda_available, _register_kernel_internal, is_pure_tensor
-
-CVCUDA_AVAILABLE = _is_cvcuda_available()
-
-if TYPE_CHECKING:
-    import cvcuda  # type: ignore[import-not-found]
-if CVCUDA_AVAILABLE:
-    cvcuda = _import_cvcuda()  # noqa: F811
+from ._utils import _get_kernel, _register_kernel_internal, is_pure_tensor
 
 
 def normalize(

torchvision/transforms/v2/functional/_utils.py

@@ -1,9 +1,13 @@
 import functools
 from collections.abc import Sequence
-from typing import Any, Callable, Optional, Union
+from typing import Any, Callable, Optional, TYPE_CHECKING, Union
 
 import torch
 from torchvision import tv_tensors
+from torchvision.transforms.functional import InterpolationMode
+
+if TYPE_CHECKING:
+    import cvcuda  # type: ignore[import-not-found]
 
 _FillType = Union[int, float, Sequence[int], Sequence[float], None]
 _FillTypeJIT = Optional[list[float]]
@@ -177,3 +181,45 @@ def _is_cvcuda_tensor(inpt: Any) -> bool:
         return isinstance(inpt, cvcuda.Tensor)
     except ImportError:
         return False
+
+
+_interpolation_mode_to_cvcuda_interp: dict[InterpolationMode | str | int, "cvcuda.Interp"] = {}
+
+
+def _populate_interpolation_mode_to_cvcuda_interp():
+    cvcuda = _import_cvcuda()
+
+    global _interpolation_mode_to_cvcuda_interp
+
+    _interpolation_mode_to_cvcuda_interp = {
+        InterpolationMode.BILINEAR: cvcuda.Interp.LINEAR,
+        "bilinear": cvcuda.Interp.LINEAR,
+        "linear": cvcuda.Interp.LINEAR,
+        2: cvcuda.Interp.LINEAR,
+        InterpolationMode.BICUBIC: cvcuda.Interp.CUBIC,
+        "bicubic": cvcuda.Interp.CUBIC,
+        3: cvcuda.Interp.CUBIC,
+        InterpolationMode.NEAREST: cvcuda.Interp.NEAREST,
+        "nearest": cvcuda.Interp.NEAREST,
+        0: cvcuda.Interp.NEAREST,
+        InterpolationMode.BOX: cvcuda.Interp.BOX,
+        "box": cvcuda.Interp.BOX,
+        4: cvcuda.Interp.BOX,
+        InterpolationMode.HAMMING: cvcuda.Interp.HAMMING,
+        "hamming": cvcuda.Interp.HAMMING,
+        5: cvcuda.Interp.HAMMING,
+        InterpolationMode.LANCZOS: cvcuda.Interp.LANCZOS,
+        "lanczos": cvcuda.Interp.LANCZOS,
+        1: cvcuda.Interp.LANCZOS,
+    }
+
+
+def _get_cvcuda_interp(interpolation: InterpolationMode | str | int) -> "cvcuda.Interp":
+    if len(_interpolation_mode_to_cvcuda_interp) == 0:
+        _populate_interpolation_mode_to_cvcuda_interp()
+
+    interp = _interpolation_mode_to_cvcuda_interp.get(interpolation)
+    if interp is None:
+        raise ValueError(f"Interpolation mode {interpolation} is not supported with CV-CUDA")
+
+    return interp