refactor Datapoint dispatch mechanism

test/common_utils.py

@@ -818,6 +818,10 @@ def make_video(size=DEFAULT_SIZE, *, num_frames=3, batch_dims=(), **kwargs):
     return datapoints.Video(make_image(size, batch_dims=(*batch_dims, num_frames), **kwargs))
 
 
+def make_video_tensor(*args, **kwargs):
+    return make_video(*args, **kwargs).as_subclass(torch.Tensor)
+
+
 def make_video_loader(
     size=DEFAULT_PORTRAIT_SPATIAL_SIZE,
     *,

test/test_transforms_v2.py

@@ -1347,8 +1347,6 @@ def test_antialias_warning():
     with pytest.warns(UserWarning, match=match):
         datapoints.Image(tensor_img).resized_crop(0, 0, 10, 10, (20, 20))
 
-    with pytest.warns(UserWarning, match=match):
-        datapoints.Video(tensor_video).resize((20, 20))
     with pytest.warns(UserWarning, match=match):
         datapoints.Video(tensor_video).resized_crop(0, 0, 10, 10, (20, 20))
 

test/test_transforms_v2_refactored.py

@@ -26,6 +26,7 @@
     make_image_tensor,
     make_segmentation_mask,
     make_video,
+    make_video_tensor,
     needs_cuda,
     set_rng_seed,
 )
@@ -39,6 +40,7 @@
 from torchvision.transforms._functional_tensor import _max_value as get_max_value
 from torchvision.transforms.functional import pil_modes_mapping
 from torchvision.transforms.v2 import functional as F
+from torchvision.transforms.v2.functional._utils import _KERNEL_REGISTRY
 
 
 @pytest.fixture(autouse=True)
@@ -177,17 +179,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 in {F.resize, F.adjust_brightness}:
+            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)
@@ -261,6 +266,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 in {F.resize, F.adjust_brightness}:
+        return
     dispatcher_signature = inspect.signature(dispatcher)
     dispatcher_params = list(dispatcher_signature.parameters.values())[1:]
 
@@ -433,6 +440,33 @@ def transform(bbox):
     return torch.stack([transform(b) for b in bounding_boxes.reshape(-1, 4).unbind()]).reshape(bounding_boxes.shape)
 
 
+@pytest.mark.parametrize(
+    ("dispatcher", "registered_datapoint_clss"),
+    [(dispatcher, set(registry.keys())) for dispatcher, registry in _KERNEL_REGISTRY.items()],
+)
+def test_exhaustive_kernel_registration(dispatcher, registered_datapoint_clss):
+    missing = {
+        datapoints.Image,
+        datapoints.BoundingBoxes,
+        datapoints.Mask,
+        datapoints.Video,
+    } - registered_datapoint_clss
+    if missing:
+        names = sorted(f"datapoints.{cls.__name__}" for cls in missing)
+        raise AssertionError(
+            "\n".join(
+                [
+                    f"The dispatcher '{dispatcher.__name__}' has no kernel registered for",
+                    "",
+                    *[f"- {name}" for name in names],
+                    "",
+                    f"If available, register the kernels with @_register_kernel_internal({dispatcher.__name__}, ...).",
+                    f"If not, register explicit no-ops with @_register_explicit_noops({', '.join(names)})",
+                ]
+            )
+        )
+
+
 class TestResize:
     INPUT_SIZE = (17, 11)
     OUTPUT_SIZES = [17, [17], (17,), [12, 13], (12, 13)]
@@ -1899,6 +1933,56 @@ def test_errors_warnings(self, make_input):
         assert out["mask"].dtype == mask_dtype
 
 
+class TestAdjustBrightness:
+    _CORRECTNESS_BRIGHTNESS_FACTORS = [0.5, 0.0, 1.0, 5.0]
+    _DEFAULT_BRIGHTNESS_FACTOR = _CORRECTNESS_BRIGHTNESS_FACTORS[0]
+
+    @pytest.mark.parametrize(
+        ("kernel", "make_input"),
+        [
+            (F.adjust_brightness_image_tensor, make_image),
+            (F.adjust_brightness_video, make_video),
+        ],
+    )
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.uint8])
+    @pytest.mark.parametrize("device", cpu_and_cuda())
+    def test_kernel(self, kernel, make_input, dtype, device):
+        check_kernel(kernel, make_input(dtype=dtype, device=device), brightness_factor=self._DEFAULT_BRIGHTNESS_FACTOR)
+
+    @pytest.mark.parametrize(
+        ("kernel", "make_input"),
+        [
+            (F.adjust_brightness_image_tensor, make_image_tensor),
+            (F.adjust_brightness_image_pil, make_image_pil),
+            (F.adjust_brightness_image_tensor, make_image),
+            (F.adjust_brightness_video, make_video),
+        ],
+    )
+    def test_dispatcher(self, kernel, make_input):
+        check_dispatcher(F.adjust_brightness, kernel, make_input(), brightness_factor=self._DEFAULT_BRIGHTNESS_FACTOR)
+
+    @pytest.mark.parametrize(
+        ("kernel", "input_type"),
+        [
+            (F.adjust_brightness_image_tensor, torch.Tensor),
+            (F.adjust_brightness_image_pil, PIL.Image.Image),
+            (F.adjust_brightness_image_tensor, datapoints.Image),
+            (F.adjust_brightness_video, datapoints.Video),
+        ],
+    )
+    def test_dispatcher_signature(self, kernel, input_type):
+        check_dispatcher_signatures_match(F.adjust_brightness, kernel=kernel, input_type=input_type)
+
+    @pytest.mark.parametrize("brightness_factor", _CORRECTNESS_BRIGHTNESS_FACTORS)
+    def test_image_correctness(self, brightness_factor):
+        image = make_image(dtype=torch.uint8, device="cpu")
+
+        actual = F.adjust_brightness(image, brightness_factor=brightness_factor)
+        expected = F.to_image_tensor(F.adjust_brightness(F.to_image_pil(image), brightness_factor=brightness_factor))
+
+        torch.testing.assert_close(actual, expected)
+
+
 class TestCutMixMixUp:
     class DummyDataset:
         def __init__(self, size, num_classes):
@@ -2036,3 +2120,92 @@ def test_labels_getter_default_heuristic(key, sample_type):
         # it takes precedence over other keys which would otherwise be a match
         d = {key: "something_else", "labels": labels}
         assert transforms._utils._find_labels_default_heuristic(d) is labels
+
+
+class TestShapeGetters:
+    @pytest.mark.parametrize(
+        ("kernel", "make_input"),
+        [
+            (F.get_dimensions_image_tensor, make_image_tensor),
+            (F.get_dimensions_image_pil, make_image_pil),
+            (F.get_dimensions_image_tensor, make_image),
+            (F.get_dimensions_video, make_video),
+        ],
+    )
+    def test_get_dimensions(self, kernel, make_input):
+        size = (10, 10)
+        color_space, num_channels = "RGB", 3
+
+        input = make_input(size, color_space=color_space)
+
+        assert kernel(input) == F.get_dimensions(input) == [num_channels, *size]
+
+    @pytest.mark.parametrize(
+        ("kernel", "make_input"),
+        [
+            (F.get_num_channels_image_tensor, make_image_tensor),
+            (F.get_num_channels_image_pil, make_image_pil),
+            (F.get_num_channels_image_tensor, make_image),
+            (F.get_num_channels_video, make_video),
+        ],
+    )
+    def test_get_num_channels(self, kernel, make_input):
+        color_space, num_channels = "RGB", 3
+
+        input = make_input(color_space=color_space)
+
+        assert kernel(input) == F.get_num_channels(input) == num_channels
+
+    @pytest.mark.parametrize(
+        ("kernel", "make_input"),
+        [
+            (F.get_size_image_tensor, make_image_tensor),
+            (F.get_size_image_pil, make_image_pil),
+            (F.get_size_image_tensor, make_image),
+            (F.get_size_bounding_boxes, make_bounding_box),
+            (F.get_size_mask, make_detection_mask),
+            (F.get_size_mask, make_segmentation_mask),
+            (F.get_size_video, make_video),
+        ],
+    )
+    def test_get_size(self, kernel, make_input):
+        size = (10, 10)
+
+        input = make_input(size)
+
+        assert kernel(input) == F.get_size(input) == list(size)
+
+    @pytest.mark.parametrize(
+        ("kernel", "make_input"),
+        [
+            (F.get_num_frames_video, make_video_tensor),
+            (F.get_num_frames_video, make_video),
+        ],
+    )
+    def test_get_num_frames(self, kernel, make_input):
+        num_frames = 4
+
+        input = make_input(num_frames=num_frames)
+
+        assert kernel(input) == F.get_num_frames(input) == num_frames
+
+    @pytest.mark.parametrize(
+        ("dispatcher", "make_input"),
+        [
+            (F.get_dimensions, make_bounding_box),
+            (F.get_dimensions, make_detection_mask),
+            (F.get_dimensions, make_segmentation_mask),
+            (F.get_num_channels, make_bounding_box),
+            (F.get_num_channels, make_detection_mask),
+            (F.get_num_channels, make_segmentation_mask),
+            (F.get_num_frames, make_image),
+            (F.get_num_frames, make_bounding_box),
+            (F.get_num_frames, make_detection_mask),
+            (F.get_num_frames, make_segmentation_mask),
+        ],
+    )
+    def test_unsupported_types(self, dispatcher, make_input):
+        input = make_input()
+
+        with pytest.raises(TypeError, match=re.escape(str(type(input)))):
+            dispatcher(input)

test/transforms_v2_dispatcher_infos.py

@@ -289,14 +289,6 @@ def fill_sequence_needs_broadcast(args_kwargs):
             skip_dispatch_datapoint,
         ],
     ),
-    DispatcherInfo(
-        F.adjust_brightness,
-        kernels={
-            datapoints.Image: F.adjust_brightness_image_tensor,
-            datapoints.Video: F.adjust_brightness_video,
-        },
-        pil_kernel_info=PILKernelInfo(F.adjust_brightness_image_pil, kernel_name="adjust_brightness_image_pil"),
-    ),
     DispatcherInfo(
         F.adjust_contrast,
         kernels={

test/transforms_v2_kernel_infos.py

@@ -1259,46 +1259,6 @@ def sample_inputs_erase_video():
     ]
 )
 
-_ADJUST_BRIGHTNESS_FACTORS = [0.1, 0.5]
-
-
-def sample_inputs_adjust_brightness_image_tensor():
-    for image_loader in make_image_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], color_spaces=("GRAY", "RGB")):
-        yield ArgsKwargs(image_loader, brightness_factor=_ADJUST_BRIGHTNESS_FACTORS[0])
-
-
-def reference_inputs_adjust_brightness_image_tensor():
-    for image_loader, brightness_factor in itertools.product(
-        make_image_loaders(color_spaces=("GRAY", "RGB"), extra_dims=[()], dtypes=[torch.uint8]),
-        _ADJUST_BRIGHTNESS_FACTORS,
-    ):
-        yield ArgsKwargs(image_loader, brightness_factor=brightness_factor)
-
-
-def sample_inputs_adjust_brightness_video():
-    for video_loader in make_video_loaders(sizes=[DEFAULT_PORTRAIT_SPATIAL_SIZE], num_frames=[3]):
-        yield ArgsKwargs(video_loader, brightness_factor=_ADJUST_BRIGHTNESS_FACTORS[0])
-
-
-KERNEL_INFOS.extend(
-    [
-        KernelInfo(
-            F.adjust_brightness_image_tensor,
-            kernel_name="adjust_brightness_image_tensor",
-            sample_inputs_fn=sample_inputs_adjust_brightness_image_tensor,
-            reference_fn=pil_reference_wrapper(F.adjust_brightness_image_pil),
-            reference_inputs_fn=reference_inputs_adjust_brightness_image_tensor,
-            float32_vs_uint8=True,
-            closeness_kwargs=float32_vs_uint8_pixel_difference(),
-        ),
-        KernelInfo(
-            F.adjust_brightness_video,
-            sample_inputs_fn=sample_inputs_adjust_brightness_video,
-        ),
-    ]
-)
-
-
 _ADJUST_CONTRAST_FACTORS = [0.1, 0.5]
 
 

torchvision/datapoints/_bounding_box.py

@@ -110,21 +110,6 @@ def vertical_flip(self) -> BoundingBoxes:
         )
         return BoundingBoxes.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",
-    ) -> BoundingBoxes:
-        output, canvas_size = self._F.resize_bounding_boxes(
-            self.as_subclass(torch.Tensor),
-            canvas_size=self.canvas_size,
-            size=size,
-            max_size=max_size,
-        )
-        return BoundingBoxes.wrap_like(self, output, canvas_size=canvas_size)
-
     def crop(self, top: int, left: int, height: int, width: int) -> BoundingBoxes:
         output, canvas_size = self._F.crop_bounding_boxes(
             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
 
@@ -228,9 +217,6 @@ def elastic(
     def rgb_to_grayscale(self, num_output_channels: int = 1) -> Datapoint:
         return self
 
-    def adjust_brightness(self, brightness_factor: float) -> Datapoint:
-        return self
-
     def adjust_saturation(self, saturation_factor: float) -> Datapoint:
         return self