refactor Datapoint dispatch mechanism

test/test_transforms_v2_refactored.py

@@ -165,17 +165,20 @@ def _check_dispatcher_dispatch(dispatcher, kernel, input, *args, **kwargs):
     preserved in doing so. For bounding boxes also checks that the format is preserved.
     """
     if isinstance(input, datapoints._datapoint.Datapoint):
-        # Due to our complex dispatch architecture for datapoints, we cannot spy on the kernel directly,
-        # but rather have to patch the `Datapoint.__F` attribute to contain the spied on kernel.
-        spy = mock.MagicMock(wraps=kernel, name=kernel.__name__)
-        with mock.patch.object(F, kernel.__name__, spy):
-            # Due to Python's name mangling, the `Datapoint.__F` attribute is only accessible from inside the class.
-            # Since that is not the case here, we need to prefix f"_{cls.__name__}"
-            # See https://docs.python.org/3/tutorial/classes.html#private-variables for details
-            with mock.patch.object(datapoints._datapoint.Datapoint, "_Datapoint__F", new=F):
-                output = dispatcher(input, *args, **kwargs)
-
-        spy.assert_called_once()
+        if dispatcher is F.resize:
+            output = dispatcher(input, *args, **kwargs)
+        else:
+            # Due to our complex dispatch architecture for datapoints, we cannot spy on the kernel directly,
+            # but rather have to patch the `Datapoint.__F` attribute to contain the spied on kernel.
+            spy = mock.MagicMock(wraps=kernel, name=kernel.__name__)
+            with mock.patch.object(F, kernel.__name__, spy):
+                # Due to Python's name mangling, the `Datapoint.__F` attribute is only accessible from inside the class.
+                # Since that is not the case here, we need to prefix f"_{cls.__name__}"
+                # See https://docs.python.org/3/tutorial/classes.html#private-variables for details
+                with mock.patch.object(datapoints._datapoint.Datapoint, "_Datapoint__F", new=F):
+                    output = dispatcher(input, *args, **kwargs)
+
+            spy.assert_called_once()
     else:
         with mock.patch(f"{dispatcher.__module__}.{kernel.__name__}", wraps=kernel) as spy:
             output = dispatcher(input, *args, **kwargs)
@@ -249,6 +252,8 @@ def _check_dispatcher_kernel_signature_match(dispatcher, *, kernel, input_type):
 
 def _check_dispatcher_datapoint_signature_match(dispatcher):
     """Checks if the signature of the dispatcher matches the corresponding method signature on the Datapoint class."""
+    if dispatcher is F.resize:
+        return
     dispatcher_signature = inspect.signature(dispatcher)
     dispatcher_params = list(dispatcher_signature.parameters.values())[1:]
 

torchvision/datapoints/_bounding_box.py

@@ -110,21 +110,6 @@ def vertical_flip(self) -> BoundingBox:
         )
         return BoundingBox.wrap_like(self, output)
 
-    def resize(  # type: ignore[override]
-        self,
-        size: List[int],
-        interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
-        max_size: Optional[int] = None,
-        antialias: Optional[Union[str, bool]] = "warn",
-    ) -> BoundingBox:
-        output, spatial_size = self._F.resize_bounding_box(
-            self.as_subclass(torch.Tensor),
-            spatial_size=self.spatial_size,
-            size=size,
-            max_size=max_size,
-        )
-        return BoundingBox.wrap_like(self, output, spatial_size=spatial_size)
-
     def crop(self, top: int, left: int, height: int, width: int) -> BoundingBox:
         output, spatial_size = self._F.crop_bounding_box(
             self.as_subclass(torch.Tensor), self.format, top=top, left=left, height=height, width=width

torchvision/datapoints/_datapoint.py

@@ -148,17 +148,6 @@ def horizontal_flip(self) -> Datapoint:
     def vertical_flip(self) -> Datapoint:
         return self
 
-    # TODO: We have to ignore override mypy error as there is torch.Tensor built-in deprecated op: Tensor.resize
-    # https://github.com/pytorch/pytorch/blob/e8727994eb7cdb2ab642749d6549bc497563aa06/torch/_tensor.py#L588-L593
-    def resize(  # type: ignore[override]
-        self,
-        size: List[int],
-        interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
-        max_size: Optional[int] = None,
-        antialias: Optional[Union[str, bool]] = "warn",
-    ) -> Datapoint:
-        return self
-
     def crop(self, top: int, left: int, height: int, width: int) -> Datapoint:
         return self
 

torchvision/datapoints/_image.py

@@ -72,18 +72,6 @@ def vertical_flip(self) -> Image:
         output = self._F.vertical_flip_image_tensor(self.as_subclass(torch.Tensor))
         return Image.wrap_like(self, output)
 
-    def resize(  # type: ignore[override]
-        self,
-        size: List[int],
-        interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
-        max_size: Optional[int] = None,
-        antialias: Optional[Union[str, bool]] = "warn",
-    ) -> Image:
-        output = self._F.resize_image_tensor(
-            self.as_subclass(torch.Tensor), size, interpolation=interpolation, max_size=max_size, antialias=antialias
-        )
-        return Image.wrap_like(self, output)
-
     def crop(self, top: int, left: int, height: int, width: int) -> Image:
         output = self._F.crop_image_tensor(self.as_subclass(torch.Tensor), top, left, height, width)
         return Image.wrap_like(self, output)

torchvision/datapoints/_mask.py

@@ -63,16 +63,6 @@ def vertical_flip(self) -> Mask:
         output = self._F.vertical_flip_mask(self.as_subclass(torch.Tensor))
         return Mask.wrap_like(self, output)
 
-    def resize(  # type: ignore[override]
-        self,
-        size: List[int],
-        interpolation: Union[InterpolationMode, int] = InterpolationMode.NEAREST,
-        max_size: Optional[int] = None,
-        antialias: Optional[Union[str, bool]] = "warn",
-    ) -> Mask:
-        output = self._F.resize_mask(self.as_subclass(torch.Tensor), size, max_size=max_size)
-        return Mask.wrap_like(self, output)
-
     def crop(self, top: int, left: int, height: int, width: int) -> Mask:
         output = self._F.crop_mask(self.as_subclass(torch.Tensor), top, left, height, width)
         return Mask.wrap_like(self, output)

torchvision/datapoints/_video.py

@@ -66,22 +66,6 @@ def vertical_flip(self) -> Video:
         output = self._F.vertical_flip_video(self.as_subclass(torch.Tensor))
         return Video.wrap_like(self, output)
 
-    def resize(  # type: ignore[override]
-        self,
-        size: List[int],
-        interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
-        max_size: Optional[int] = None,
-        antialias: Optional[Union[str, bool]] = "warn",
-    ) -> Video:
-        output = self._F.resize_video(
-            self.as_subclass(torch.Tensor),
-            size,
-            interpolation=interpolation,
-            max_size=max_size,
-            antialias=antialias,
-        )
-        return Video.wrap_like(self, output)
-
     def crop(self, top: int, left: int, height: int, width: int) -> Video:
         output = self._F.crop_video(self.as_subclass(torch.Tensor), top, left, height, width)
         return Video.wrap_like(self, output)

torchvision/transforms/v2/functional/__init__.py

@@ -1,6 +1,6 @@
 from torchvision.transforms import InterpolationMode  # usort: skip
 
-from ._utils import is_simple_tensor  # usort: skip
+from ._utils import is_simple_tensor, register_kernel  # usort: skip
 
 from ._meta import (
     clamp_bounding_box,

torchvision/transforms/v2/functional/_geometry.py

@@ -25,7 +25,7 @@
 
 from ._meta import clamp_bounding_box, convert_format_bounding_box, get_spatial_size_image_pil
 
-from ._utils import is_simple_tensor
+from ._utils import _get_kernel, is_simple_tensor, register_kernel
 
 
 def _check_interpolation(interpolation: Union[InterpolationMode, int]) -> InterpolationMode:
@@ -158,6 +158,32 @@ def _compute_resized_output_size(
     return __compute_resized_output_size(spatial_size, size=size, max_size=max_size)
 
 
+def resize(
+    inpt: datapoints._InputTypeJIT,
+    size: List[int],
+    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
+    max_size: Optional[int] = None,
+    antialias: Optional[Union[str, bool]] = "warn",
+) -> datapoints._InputTypeJIT:
+    if not torch.jit.is_scripting():
+        _log_api_usage_once(resize)
+    if torch.jit.is_scripting() or is_simple_tensor(inpt):
+        return resize_image_tensor(inpt, size, interpolation=interpolation, max_size=max_size, antialias=antialias)
+    elif isinstance(inpt, datapoints._datapoint.Datapoint):
+        kernel = _get_kernel(resize, type(inpt))
+        return kernel(inpt, size, interpolation=interpolation, max_size=max_size, antialias=antialias)
+    elif isinstance(inpt, PIL.Image.Image):
+        if antialias is False:
+            warnings.warn("Anti-alias option is always applied for PIL Image input. Argument antialias is ignored.")
+        return resize_image_pil(inpt, size, interpolation=interpolation, max_size=max_size)
+    else:
+        raise TypeError(
+            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
+            f"but got {type(inpt)} instead."
+        )
+
+
+@register_kernel(resize, datapoints.Image)
 def resize_image_tensor(
     image: torch.Tensor,
     size: List[int],
@@ -274,6 +300,11 @@ def resize_mask(mask: torch.Tensor, size: List[int], max_size: Optional[int] = N
     return output
 
 
+@register_kernel(resize, datapoints.Mask)
+def _resize_mask_dispatch(mask: torch.Tensor, size: List[int], max_size: Optional[int] = None, **kwargs):
+    return resize_mask(mask, size, max_size=max_size)
+
+
 def resize_bounding_box(
     bounding_box: torch.Tensor, spatial_size: Tuple[int, int], size: List[int], max_size: Optional[int] = None
 ) -> Tuple[torch.Tensor, Tuple[int, int]]:
@@ -292,6 +323,17 @@ def resize_bounding_box(
     )
 
 
+@register_kernel(resize, datapoints.BoundingBox, datapoint_wrapping=False)
+def _resize_bounding_box_dispatch(
+    bounding_box: datapoints.BoundingBox, size: List[int], max_size: Optional[int] = None, **kwargs
+):
+    output, spatial_size = resize_bounding_box(
+        bounding_box.as_subclass(torch.Tensor), bounding_box.spatial_size, size, max_size=max_size
+    )
+    return datapoints.BoundingBox.wrap_like(bounding_box, output, spatial_size=spatial_size)
+
+
+@register_kernel(resize, datapoints.Video)
 def resize_video(
     video: torch.Tensor,
     size: List[int],
@@ -302,30 +344,6 @@ def resize_video(
     return resize_image_tensor(video, size=size, interpolation=interpolation, max_size=max_size, antialias=antialias)
 
 
-def resize(
-    inpt: datapoints._InputTypeJIT,
-    size: List[int],
-    interpolation: Union[InterpolationMode, int] = InterpolationMode.BILINEAR,
-    max_size: Optional[int] = None,
-    antialias: Optional[Union[str, bool]] = "warn",
-) -> datapoints._InputTypeJIT:
-    if not torch.jit.is_scripting():
-        _log_api_usage_once(resize)
-    if torch.jit.is_scripting() or is_simple_tensor(inpt):
-        return resize_image_tensor(inpt, size, interpolation=interpolation, max_size=max_size, antialias=antialias)
-    elif isinstance(inpt, datapoints._datapoint.Datapoint):
-        return inpt.resize(size, interpolation=interpolation, max_size=max_size, antialias=antialias)
-    elif isinstance(inpt, PIL.Image.Image):
-        if antialias is False:
-            warnings.warn("Anti-alias option is always applied for PIL Image input. Argument antialias is ignored.")
-        return resize_image_pil(inpt, size, interpolation=interpolation, max_size=max_size)
-    else:
-        raise TypeError(
-            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
-            f"but got {type(inpt)} instead."
-        )
-
-
 def _affine_parse_args(
     angle: Union[int, float],
     translate: List[float],

torchvision/transforms/v2/functional/_utils.py

@@ -1,3 +1,4 @@
+import functools
 from typing import Any
 
 import torch
@@ -6,3 +7,25 @@
 
 def is_simple_tensor(inpt: Any) -> bool:
     return isinstance(inpt, torch.Tensor) and not isinstance(inpt, Datapoint)
+
+
+_KERNEL_REGISTRY = {}
+
+
+def register_kernel(dispatcher, datapoint_cls, *, datapoint_wrapping=True):
+    def datapoint_wrapper(kernel):
+        @functools.wraps(kernel)
+        def wrapper(inpt, *args, **kwargs):
+            return type(inpt).wrap_like(inpt, kernel(inpt.as_subclass(torch.Tensor), *args, **kwargs))
+
+        return wrapper
+
+    def decorator(kernel):
+        _KERNEL_REGISTRY[(dispatcher, datapoint_cls)] = datapoint_wrapper(kernel) if datapoint_wrapping else kernel
+        return kernel
+
+    return decorator
+
+
+def _get_kernel(dispatcher, datapoint_cls):
+    return _KERNEL_REGISTRY[(dispatcher, datapoint_cls)]