Updated conditions for _BaseMixupCutmix

Author: vfdev-5

test/test_prototype_transforms.py

@@ -7,8 +7,6 @@
     make_images,
     make_bounding_boxes,
     make_one_hot_labels,
-    make_segmentation_masks,
-    make_label,
 )
 from torchvision.prototype import transforms, features
 from torchvision.transforms.functional import to_pil_image, pil_to_tensor
@@ -108,20 +106,6 @@ def test_common(self, transform, input):
     def test_mixup_cutmix(self, transform, input):
         transform(input)
 
-    @pytest.mark.parametrize("transform", [transforms.RandomMixup(alpha=1.0), transforms.RandomCutmix(alpha=1.0)])
-    def test_mixup_cutmix_assertions(self, transform):
-        for bbox in make_bounding_boxes():
-            with pytest.raises(TypeError, match="does not support"):
-                transform(bbox)
-            break
-        for mask in make_segmentation_masks():
-            with pytest.raises(TypeError, match="does not support"):
-                transform(mask)
-            break
-        label = make_label()
-        with pytest.raises(TypeError, match="does not support"):
-            transform(label)
-
     @parametrize(
         [
             (

torchvision/prototype/transforms/_augment.py

@@ -9,7 +9,7 @@
 from torchvision.prototype.transforms import Transform, functional as F
 
 from ._transform import _RandomApplyTransform
-from ._utils import query_image, get_image_dimensions, has_any
+from ._utils import query_image, get_image_dimensions, has_any, has_all
 
 
 class RandomErasing(_RandomApplyTransform):
@@ -106,6 +106,8 @@ def __init__(self, *, alpha: float) -> None:
 
     def forward(self, *inpts: Any) -> Any:
         sample = inpts if len(inpts) > 1 else inpts[0]
+        if not has_all(sample, features.Image, features.OneHotLabel):
+            raise TypeError(f"{type(self).__name__}() is only defined for Image's *and* OneHotLabel's.")
         if has_any(sample, features.BoundingBox, features.SegmentationMask, features.Label):
             raise TypeError(
                 f"{type(self).__name__}() does not support bounding boxes, segmentation masks and plain labels."