Transforms without dispatcher

test/test_prototype_transforms.py

@@ -1,6 +1,5 @@
 import itertools
 
-import PIL.Image
 import pytest
 import torch
 from test_prototype_transforms_kernels import make_images, make_bounding_boxes, make_one_hot_labels
@@ -25,15 +24,6 @@ def make_vanilla_tensor_bounding_boxes(*args, **kwargs):
         yield bounding_box.data
 
 
-INPUT_CREATIONS_FNS = {
-    features.Image: make_images,
-    features.BoundingBox: make_bounding_boxes,
-    features.OneHotLabel: make_one_hot_labels,
-    torch.Tensor: make_vanilla_tensor_images,
-    PIL.Image.Image: make_pil_images,
-}
-
-
 def parametrize(transforms_with_inputs):
     return pytest.mark.parametrize(
         ("transform", "input"),
@@ -52,15 +42,21 @@ def parametrize(transforms_with_inputs):
 def parametrize_from_transforms(*transforms):
     transforms_with_inputs = []
     for transform in transforms:
-        dispatcher = transform._DISPATCHER
-        if dispatcher is None:
-            continue
-
-        for type_ in dispatcher._kernels:
+        for creation_fn in [
+            make_images,
+            make_bounding_boxes,
+            make_one_hot_labels,
+            make_vanilla_tensor_images,
+            make_pil_images,
+        ]:
+            inputs = list(creation_fn())
             try:
-                inputs = INPUT_CREATIONS_FNS[type_]()
-            except KeyError:
+                output = transform(inputs[0])
+            except Exception:
                 continue
+            else:
+                if output is inputs[0]:
+                    continue
 
             transforms_with_inputs.append((transform, inputs))
 
@@ -69,7 +65,7 @@ def parametrize_from_transforms(*transforms):
 
 class TestSmoke:
     @parametrize_from_transforms(
-        transforms.RandomErasing(),
+        transforms.RandomErasing(p=1.0),
         transforms.HorizontalFlip(),
         transforms.Resize([16, 16]),
         transforms.CenterCrop([16, 16]),
@@ -141,35 +137,6 @@ def test_auto_augment(self, transform, input):
     def test_normalize(self, transform, input):
         transform(input)
 
-    @parametrize(
-        [
-            (
-                transforms.ConvertColorSpace("grayscale"),
-                itertools.chain(
-                    make_images(),
-                    make_vanilla_tensor_images(color_spaces=["rgb"]),
-                    make_pil_images(color_spaces=["rgb"]),
-                ),
-            )
-        ]
-    )
-    def test_convert_bounding_color_space(self, transform, input):
-        transform(input)
-
-    @parametrize(
-        [
-            (
-                transforms.ConvertBoundingBoxFormat("xyxy", old_format="xywh"),
-                itertools.chain(
-                    make_bounding_boxes(),
-                    make_vanilla_tensor_bounding_boxes(formats=["xywh"]),
-                ),
-            )
-        ]
-    )
-    def test_convert_bounding_box_format(self, transform, input):
-        transform(input)
-
     @parametrize(
         [
             (

torchvision/prototype/transforms/__init__.py

@@ -1,6 +1,7 @@
-from torchvision.transforms import AutoAugmentPolicy, InterpolationMode  # usort: skip
+from torchvision.transforms import InterpolationMode, AutoAugmentPolicy  # usort: skip
+
 from . import kernels  # usort: skip
-from . import functional  # usort: skip
+
 from ._transform import Transform  # usort: skip
 
 from ._augment import RandomErasing, RandomMixup, RandomCutmix

torchvision/prototype/transforms/_augment.py

@@ -6,15 +6,13 @@
 import PIL.Image
 import torch
 from torchvision.prototype import features
-from torchvision.prototype.transforms import Transform, functional as F
+from torchvision.prototype.transforms import Transform, kernels as K
+from torchvision.transforms import functional as _F
 
-from ._utils import query_image
+from ._utils import query_image, get_image_size, get_image_num_channels
 
 
 class RandomErasing(Transform):
-    _DISPATCHER = F.erase
-    _FAIL_TYPES = {PIL.Image.Image, features.BoundingBox, features.SegmentationMask}
-
     def __init__(
         self,
         p: float = 0.5,
@@ -45,8 +43,8 @@ def __init__(
 
     def _get_params(self, sample: Any) -> Dict[str, Any]:
         image = query_image(sample)
-        img_h, img_w = F.get_image_size(image)
-        img_c = F.get_image_num_channels(image)
+        img_c = get_image_num_channels(image)
+        img_h, img_w = get_image_size(image)
 
         if isinstance(self.value, (int, float)):
             value = [self.value]
@@ -59,7 +57,8 @@ def _get_params(self, sample: Any) -> Dict[str, Any]:
 
         if value is not None and not (len(value) in (1, img_c)):
             raise ValueError(
-                f"If value is a sequence, it should have either a single value or {img_c} (number of input channels)"
+                "If value is a sequence, it should have either a single value or "
+                f"{image.shape[-3]} (number of input channels)"
             )
 
         area = img_h * img_w
@@ -93,16 +92,23 @@ def _get_params(self, sample: Any) -> Dict[str, Any]:
         return dict(zip("ijhwv", (i, j, h, w, v)))
 
     def _transform(self, input: Any, params: Dict[str, Any]) -> Any:
-        if torch.rand(1) >= self.p:
+        if type(input) is torch.Tensor:
+            return _F.erase(input, **params)
+        elif type(input) is features.Image:
+            return features.Image.new_like(input, K.erase_image(input, **params))
+        elif type(input) in {features.BoundingBox, features.SegmentationMask} or isinstance(input, PIL.Image.Image):
+            raise TypeError(f"{type(input)} is not supported by {type(self).__name__}()")
+        else:
             return input
 
-        return super()._transform(input, params)
+    def forward(self, *inputs: Any) -> Any:
+        if torch.rand(1) >= self.p:
+            return inputs if len(inputs) > 1 else inputs[0]
+
+        return super().forward(*inputs)
 
 
 class RandomMixup(Transform):
-    _DISPATCHER = F.mixup
-    _FAIL_TYPES = {features.BoundingBox, features.SegmentationMask}
-
     def __init__(self, *, alpha: float) -> None:
         super().__init__()
         self.alpha = alpha
@@ -111,11 +117,20 @@ def __init__(self, *, alpha: float) -> None:
     def _get_params(self, sample: Any) -> Dict[str, Any]:
         return dict(lam=float(self._dist.sample(())))
 
+    def _transform(self, input: Any, params: Dict[str, Any]) -> Any:
+        if type(input) is features.Image:
+            output = K.mixup_image(input, **params)
+            return features.Image.new_like(input, output)
+        elif type(input) is features.OneHotLabel:
+            output = K.mixup_one_hot_label(input, **params)
+            return features.OneHotLabel.new_like(input, output)
+        elif type(input) in {torch.Tensor, features.BoundingBox, features.SegmentationMask}:
+            raise TypeError(f"{type(input)} is not supported by {type(self).__name__}()")
+        else:
+            return input
 
-class RandomCutmix(Transform):
-    _DISPATCHER = F.cutmix
-    _FAIL_TYPES = {features.BoundingBox, features.SegmentationMask}
 
+class RandomCutmix(Transform):
     def __init__(self, *, alpha: float) -> None:
         super().__init__()
         self.alpha = alpha
@@ -125,7 +140,7 @@ def _get_params(self, sample: Any) -> Dict[str, Any]:
         lam = float(self._dist.sample(()))
 
         image = query_image(sample)
-        H, W = F.get_image_size(image)
+        H, W = get_image_size(image)
 
         r_x = torch.randint(W, ())
         r_y = torch.randint(H, ())
@@ -143,3 +158,15 @@ def _get_params(self, sample: Any) -> Dict[str, Any]:
         lam_adjusted = float(1.0 - (x2 - x1) * (y2 - y1) / (W * H))
 
         return dict(box=box, lam_adjusted=lam_adjusted)
+
+    def _transform(self, input: Any, params: Dict[str, Any]) -> Any:
+        if type(input) is features.Image:
+            output = K.cutmix_image(input, box=params["box"])
+            return features.Image.new_like(input, output)
+        elif type(input) is features.OneHotLabel:
+            output = K.cutmix_one_hot_label(input, lam_adjusted=params["lam_adjusted"])
+            return features.OneHotLabel.new_like(input, output)
+        elif type(input) in {torch.Tensor, features.BoundingBox, features.SegmentationMask}:
+            raise TypeError(f"{type(input)} is not supported by {type(self).__name__}()")
+        else:
+            return input