refactor Datapoint dispatch mechanism

test/common_utils.py

@@ -829,6 +829,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/datasets_utils.py

@@ -567,7 +567,7 @@ def test_transforms(self, config):
 
     @test_all_configs
     def test_transforms_v2_wrapper(self, config):
-        from torchvision.datapoints._datapoint import Datapoint
+        from torchvision import datapoints
         from torchvision.datasets import wrap_dataset_for_transforms_v2
 
         try:
@@ -588,7 +588,9 @@ def test_transforms_v2_wrapper(self, config):
                     assert len(wrapped_dataset) == info["num_examples"]
 
                     wrapped_sample = wrapped_dataset[0]
-                    assert tree_any(lambda item: isinstance(item, (Datapoint, PIL.Image.Image)), wrapped_sample)
+                    assert tree_any(
+                        lambda item: isinstance(item, (datapoints.Datapoint, PIL.Image.Image)), wrapped_sample
+                    )
         except TypeError as error:
             msg = f"No wrapper exists for dataset class {type(dataset).__name__}"
             if str(error).startswith(msg):

test/test_transforms_v2.py

@@ -1346,8 +1346,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_functional.py

@@ -3,8 +3,6 @@
 import os
 import re
 
-from typing import get_type_hints
-
 import numpy as np
 import PIL.Image
 import pytest
@@ -417,22 +415,6 @@ def test_pil_output_type(self, info, args_kwargs):
 
         assert isinstance(output, PIL.Image.Image)
 
-    @make_info_args_kwargs_parametrization(
-        DISPATCHER_INFOS,
-        args_kwargs_fn=lambda info: info.sample_inputs(),
-    )
-    def test_dispatch_datapoint(self, info, args_kwargs, spy_on):
-        (datapoint, *other_args), kwargs = args_kwargs.load()
-
-        method_name = info.id
-        method = getattr(datapoint, method_name)
-        datapoint_type = type(datapoint)
-        spy = spy_on(method, module=datapoint_type.__module__, name=f"{datapoint_type.__name__}.{method_name}")
-
-        info.dispatcher(datapoint, *other_args, **kwargs)
-
-        spy.assert_called_once()
-
     @make_info_args_kwargs_parametrization(
         DISPATCHER_INFOS,
         args_kwargs_fn=lambda info: info.sample_inputs(),
@@ -462,9 +444,12 @@ def test_dispatcher_kernel_signatures_consistency(self, dispatcher_info, datapoi
         kernel_params = list(kernel_signature.parameters.values())[1:]
 
         # We filter out metadata that is implicitly passed to the dispatcher through the input datapoint, but has to be
-        # explicit passed to the kernel.
-        datapoint_type_metadata = datapoint_type.__annotations__.keys()
-        kernel_params = [param for param in kernel_params if param.name not in datapoint_type_metadata]
+        # explicitly passed to the kernel.
+        input_type = {v: k for k, v in dispatcher_info.kernels.items()}.get(kernel_info.kernel)
+        explicit_metadata = {
+            datapoints.BoundingBoxes: {"format", "canvas_size"},
+        }
+        kernel_params = [param for param in kernel_params if param.name not in explicit_metadata.get(input_type, set())]
 
         dispatcher_params = iter(dispatcher_params)
         for dispatcher_param, kernel_param in zip(dispatcher_params, kernel_params):
@@ -481,28 +466,6 @@ def test_dispatcher_kernel_signatures_consistency(self, dispatcher_info, datapoi
 
             assert dispatcher_param == kernel_param
 
-    @pytest.mark.parametrize("info", DISPATCHER_INFOS, ids=lambda info: info.id)
-    def test_dispatcher_datapoint_signatures_consistency(self, info):
-        try:
-            datapoint_method = getattr(datapoints._datapoint.Datapoint, info.id)
-        except AttributeError:
-            pytest.skip("Dispatcher doesn't support arbitrary datapoint dispatch.")
-
-        dispatcher_signature = inspect.signature(info.dispatcher)
-        dispatcher_params = list(dispatcher_signature.parameters.values())[1:]
-
-        datapoint_signature = inspect.signature(datapoint_method)
-        datapoint_params = list(datapoint_signature.parameters.values())[1:]
-
-        # Because we use `from __future__ import annotations` inside the module where `datapoints._datapoint` is
-        # defined, the annotations are stored as strings. This makes them concrete again, so they can be compared to the
-        # natively concrete dispatcher annotations.
-        datapoint_annotations = get_type_hints(datapoint_method)
-        for param in datapoint_params:
-            param._annotation = datapoint_annotations[param.name]
-
-        assert dispatcher_params == datapoint_params
-
     @pytest.mark.parametrize("info", DISPATCHER_INFOS, ids=lambda info: info.id)
     def test_unkown_type(self, info):
         unkown_input = object()