NEW Feature: Mixup transform for Object Detection

test/test_prototype_transforms.py

@@ -26,6 +26,7 @@
 from torchvision.ops.boxes import box_iou
 from torchvision.prototype import datapoints, transforms
 from torchvision.prototype.transforms.utils import check_type
+
 from torchvision.transforms.functional import InterpolationMode, pil_to_tensor, to_pil_image
 
 BATCH_EXTRA_DIMS = [extra_dims for extra_dims in DEFAULT_EXTRA_DIMS if extra_dims]
@@ -1436,63 +1437,6 @@ def create_fake_image(self, mocker, image_type):
             return PIL.Image.new("RGB", (32, 32), 123)
         return mocker.MagicMock(spec=image_type)
 
-    def test__extract_image_targets_assertion(self, mocker):
-        transform = transforms.SimpleCopyPaste()
-
-        flat_sample = [
-            # images, batch size = 2
-            self.create_fake_image(mocker, datapoints.Image),
-            # labels, bboxes, masks
-            mocker.MagicMock(spec=datapoints.Label),
-            mocker.MagicMock(spec=datapoints.BoundingBox),
-            mocker.MagicMock(spec=datapoints.Mask),
-            # labels, bboxes, masks
-            mocker.MagicMock(spec=datapoints.BoundingBox),
-            mocker.MagicMock(spec=datapoints.Mask),
-        ]
-
-        with pytest.raises(TypeError, match="requires input sample to contain equal sized list of Images"):
-            transform._extract_image_targets(flat_sample)
-
-    @pytest.mark.parametrize("image_type", [datapoints.Image, PIL.Image.Image, torch.Tensor])
-    @pytest.mark.parametrize("label_type", [datapoints.Label, datapoints.OneHotLabel])
-    def test__extract_image_targets(self, image_type, label_type, mocker):
-        transform = transforms.SimpleCopyPaste()
-
-        flat_sample = [
-            # images, batch size = 2
-            self.create_fake_image(mocker, image_type),
-            self.create_fake_image(mocker, image_type),
-            # labels, bboxes, masks
-            mocker.MagicMock(spec=label_type),
-            mocker.MagicMock(spec=datapoints.BoundingBox),
-            mocker.MagicMock(spec=datapoints.Mask),
-            # labels, bboxes, masks
-            mocker.MagicMock(spec=label_type),
-            mocker.MagicMock(spec=datapoints.BoundingBox),
-            mocker.MagicMock(spec=datapoints.Mask),
-        ]
-
-        images, targets = transform._extract_image_targets(flat_sample)
-
-        assert len(images) == len(targets) == 2
-        if image_type == PIL.Image.Image:
-            torch.testing.assert_close(images[0], pil_to_tensor(flat_sample[0]))
-            torch.testing.assert_close(images[1], pil_to_tensor(flat_sample[1]))
-        else:
-            assert images[0] == flat_sample[0]
-            assert images[1] == flat_sample[1]
-
-        for target in targets:
-            for key, type_ in [
-                ("boxes", datapoints.BoundingBox),
-                ("masks", datapoints.Mask),
-                ("labels", label_type),
-            ]:
-                assert key in target
-                assert isinstance(target[key], type_)
-                assert target[key] in flat_sample
-
     @pytest.mark.parametrize("label_type", [datapoints.Label, datapoints.OneHotLabel])
     def test__copy_paste(self, label_type):
         image = 2 * torch.ones(3, 32, 32)
@@ -1924,3 +1868,49 @@ def test__transform(self, inpt):
         assert type(output) is type(inpt)
         assert output.shape[-4] == num_samples
         assert output.dtype == inpt.dtype
+
+
+class TestMixupDetection:
+    def create_fake_image(self, mocker, image_type):
+        if image_type == PIL.Image.Image:
+            return PIL.Image.new("RGB", (32, 32), 123)
+        return mocker.MagicMock(spec=image_type)
+
+    def test__mixup(self):
+        image1 = 2 * torch.ones(3, 32, 64)
+        target_1 = {
+            "boxes": datapoints.BoundingBox(
+                torch.tensor([[0.0, 0.0, 10.0, 10.0], [20.0, 20.0, 30.0, 30.0]]),
+                format="XYXY",
+                spatial_size=(32, 64),
+            ),
+            "labels": datapoints.Label(torch.tensor([1, 2])),
+        }
+
+        image2 = 10 * torch.ones(3, 64, 32)
+        target_2 = {
+            "boxes": datapoints.BoundingBox(
+                torch.tensor([[10.0, 0.0, 20.0, 20.0], [10.0, 20.0, 30.0, 30.0]]),
+                format="XYXY",
+                spatial_size=(64, 32),
+            ),
+            "labels": datapoints.Label(torch.tensor([2, 3])),
+        }
+
+        sample_1 = {
+            "image": image1,
+            "boxes": target_1["boxes"],
+            "labels": target_1["labels"],
+        }
+        sample_2 = {
+            "image": image2,
+            "boxes": target_2["boxes"],
+            "labels": target_2["labels"],
+        }
+
+        transform = transforms.MixupDetection()
+        output = transform._mixup(sample_1, sample_2, 0.5)
+        assert output["image"].shape == (3, 64, 64)
+        assert output["boxes"].spatial_size == (64, 64)
+        assert len(output["boxes"]) == 4
+        assert len(output["labels"]) == 4

torchvision/prototype/transforms/__init__.py

@@ -5,7 +5,7 @@
 from ._transform import Transform  # usort: skip
 from ._presets import StereoMatching  # usort: skip
 
-from ._augment import RandomCutmix, RandomErasing, RandomMixup, SimpleCopyPaste
+from ._augment import MixupDetection, RandomCutmix, RandomErasing, RandomMixup, SimpleCopyPaste
 from ._auto_augment import AugMix, AutoAugment, RandAugment, TrivialAugmentWide
 from ._color import (
     ColorJitter,

torchvision/prototype/transforms/_augment.py

@@ -4,14 +4,13 @@
 from typing import Any, cast, Dict, List, Optional, Tuple, Union
 
 import PIL.Image
-import torch
-from torch.utils._pytree import tree_flatten, tree_unflatten
 
+import torch
 from torchvision.ops import masks_to_boxes
 from torchvision.prototype import datapoints
 from torchvision.prototype.transforms import functional as F, InterpolationMode, Transform
 
-from ._transform import _RandomApplyTransform
+from ._transform import _DetectionBatchTransform, _RandomApplyTransform
 from .utils import has_any, is_simple_tensor, query_chw, query_spatial_size
 
 
@@ -214,7 +213,6 @@ def _copy_paste(
         resize_interpolation: F.InterpolationMode,
         antialias: Optional[bool],
     ) -> Tuple[datapoints.TensorImageType, Dict[str, Any]]:
-
         paste_masks = paste_target["masks"].wrap_like(paste_target["masks"], paste_target["masks"][random_selection])
         paste_boxes = paste_target["boxes"].wrap_like(paste_target["boxes"], paste_target["boxes"][random_selection])
         paste_labels = paste_target["labels"].wrap_like(
@@ -241,7 +239,7 @@ def _copy_paste(
 
         inverse_paste_alpha_mask = paste_alpha_mask.logical_not()
         # Copy-paste images:
-        image = image.mul(inverse_paste_alpha_mask).add_(paste_image.mul(paste_alpha_mask))
+        out_image = image.mul(inverse_paste_alpha_mask).add_(paste_image.mul(paste_alpha_mask))
 
         # Copy-paste masks:
         masks = masks * inverse_paste_alpha_mask
@@ -281,69 +279,15 @@ def _copy_paste(
             out_target["masks"] = out_target["masks"][valid_targets]
             out_target["labels"] = out_target["labels"][valid_targets]
 
-        return image, out_target
-
-    def _extract_image_targets(
-        self, flat_sample: List[Any]
-    ) -> Tuple[List[datapoints.TensorImageType], List[Dict[str, Any]]]:
-        # fetch all images, bboxes, masks and labels from unstructured input
-        # with List[image], List[BoundingBox], List[Mask], List[Label]
-        images, bboxes, masks, labels = [], [], [], []
-        for obj in flat_sample:
-            if isinstance(obj, datapoints.Image) or is_simple_tensor(obj):
-                images.append(obj)
-            elif isinstance(obj, PIL.Image.Image):
-                images.append(F.to_image_tensor(obj))
-            elif isinstance(obj, datapoints.BoundingBox):
-                bboxes.append(obj)
-            elif isinstance(obj, datapoints.Mask):
-                masks.append(obj)
-            elif isinstance(obj, (datapoints.Label, datapoints.OneHotLabel)):
-                labels.append(obj)
-
-        if not (len(images) == len(bboxes) == len(masks) == len(labels)):
-            raise TypeError(
-                f"{type(self).__name__}() requires input sample to contain equal sized list of Images, "
-                "BoundingBoxes, Masks and Labels or OneHotLabels."
-            )
-
-        targets = []
-        for bbox, mask, label in zip(bboxes, masks, labels):
-            targets.append({"boxes": bbox, "masks": mask, "labels": label})
-
-        return images, targets
-
-    def _insert_outputs(
-        self,
-        flat_sample: List[Any],
-        output_images: List[datapoints.TensorImageType],
-        output_targets: List[Dict[str, Any]],
-    ) -> None:
-        c0, c1, c2, c3 = 0, 0, 0, 0
-        for i, obj in enumerate(flat_sample):
-            if isinstance(obj, datapoints.Image):
-                flat_sample[i] = datapoints.Image.wrap_like(obj, output_images[c0])
-                c0 += 1
-            elif isinstance(obj, PIL.Image.Image):
-                flat_sample[i] = F.to_image_pil(output_images[c0])
-                c0 += 1
-            elif is_simple_tensor(obj):
-                flat_sample[i] = output_images[c0]
-                c0 += 1
-            elif isinstance(obj, datapoints.BoundingBox):
-                flat_sample[i] = datapoints.BoundingBox.wrap_like(obj, output_targets[c1]["boxes"])
-                c1 += 1
-            elif isinstance(obj, datapoints.Mask):
-                flat_sample[i] = datapoints.Mask.wrap_like(obj, output_targets[c2]["masks"])
-                c2 += 1
-            elif isinstance(obj, (datapoints.Label, datapoints.OneHotLabel)):
-                flat_sample[i] = obj.wrap_like(obj, output_targets[c3]["labels"])  # type: ignore[arg-type]
-                c3 += 1
+        return out_image, out_target
 
     def forward(self, *inputs: Any) -> Any:
-        flat_inputs, spec = tree_flatten(inputs if len(inputs) > 1 else inputs[0])
-
-        images, targets = self._extract_image_targets(flat_inputs)
+        flat_batch_with_spec, images, targets = flatten_and_extract_data(
+            inputs,
+            boxes=(datapoints.BoundingBox,),
+            masks=(datapoints.Mask,),
+            labels=(datapoints.Label, datapoints.OneHotLabel),
+        )
 
         # images = [t1, t2, ..., tN]
         # Let's define paste_images as shifted list of input images
@@ -380,7 +324,73 @@ def forward(self, *inputs: Any) -> Any:
             output_images.append(output_image)
             output_targets.append(output_target)
 
-        # Insert updated images and targets into input flat_sample
-        self._insert_outputs(flat_inputs, output_images, output_targets)
+        return unflatten_and_insert_data(flat_batch_with_spec, output_images, output_targets)
+
+
+class MixupDetection(_DetectionBatchTransform):
+    _transformed_types = (is_simple_tensor, datapoints.Image, PIL.Image)
+
+    def __init__(
+        self,
+        *,
+        alpha: float = 1.5,
+    ) -> None:
+        super().__init__()
+        self._dist = torch.distributions.Beta(torch.tensor([alpha]), torch.tensor([alpha]))
+
+    def _check_inputs(self, flat_inputs: List[Any]) -> None:
+        if has_any(flat_inputs, datapoints.Mask, datapoints.Video):
+            raise TypeError(f"{type(self).__name__}() is only supported for images and bounding boxes.")
+
+    def _get_params(self, flat_inputs: List[Any]) -> Dict[str, Any]:
+        return dict(ratio=float(self._dist.sample()))
+
+    def forward(self, *inputs: Any) -> Any:
+        flat_batch_with_spec, batch = self._flatten_and_extract_data(
+            inputs,
+            image=(datapoints.Image, PIL.Image.Image, is_simple_tensor),
+            boxes=(datapoints.BoundingBox,),
+            labels=(datapoints.Label, datapoints.OneHotLabel),
+        )
+        self._check_inputs(flat_batch_with_spec[0])
+
+        batch = self._to_image_tensor(batch)
+
+        batch_output = [
+            self._mixup(sample, sample_rolled, self._get_params([])["ratio"])
+            for sample, sample_rolled in zip(batch, batch[-1:] + batch[:-1])
+        ]
+
+        return self._unflatten_and_insert_data(flat_batch_with_spec, batch_output)
+
+    def _mixup(self, sample_1: Dict[str, Any], sample_2: Dict[str, Any], ratio: float) -> Dict[str, Any]:
+        if ratio >= 1.0:
+            return sample_1
+        elif ratio == 0.0:
+            return sample_2
+
+        h_1, w_1 = sample_1["image"].shape[-2:]
+        h_2, w_2 = sample_2["image"].shape[-2:]
+        h_mixup = max(h_1, h_2)
+        w_mixup = max(w_1, w_2)
+
+        # TODO: add the option to fill this with something else than 0
+        dtype = sample_1["image"].dtype if sample_1["image"].is_floating_point() else torch.float32
+        mix_image = F.pad_image_tensor(
+            sample_1["image"].to(dtype), padding=[0, 0, w_mixup - w_1, h_mixup - h_1], fill=None
+        ).mul_(ratio)
+        mix_image[..., :h_2, :w_2] = sample_2["image"] * (1.0 - ratio)
+        mix_image = mix_image.to(sample_1["image"])
+
+        mix_boxes = datapoints.BoundingBox.wrap_like(
+            sample_1["boxes"],
+            torch.cat([sample_1["boxes"], sample_2["boxes"]], dim=-2),
+            spatial_size=(h_mixup, w_mixup),
+        )
+
+        mix_labels = datapoints.Label.wrap_like(
+            sample_1["labels"],
+            torch.cat([sample_1["labels"], sample_2["labels"]], dim=-1),
+        )
 
-        return tree_unflatten(flat_inputs, spec)
+        return dict(image=mix_image, boxes=mix_boxes, labels=mix_labels)