NEW Feature: Mixup transform for Object Detection

test/test_prototype_transforms.py

@@ -24,7 +24,7 @@
 from torchvision.ops.boxes import box_iou
 from torchvision.prototype import features, transforms
 from torchvision.prototype.transforms._utils import _isinstance
-from torchvision.transforms.functional import InterpolationMode, pil_to_tensor, to_pil_image
+from torchvision.transforms.functional import InterpolationMode, pil_to_tensor, to_pil_image, to_tensor
 
 BATCH_EXTRA_DIMS = [extra_dims for extra_dims in DEFAULT_EXTRA_DIMS if extra_dims]
 
@@ -1918,3 +1918,80 @@ 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__extract_image_targets_assertion(self, mocker):
+        transform = transforms.MixupDetection()
+
+        flat_sample = [
+            # images, batch size = 2
+            self.create_fake_image(mocker, features.Image),
+            # labels, bboxes, masks
+            mocker.MagicMock(spec=features.Label),
+            mocker.MagicMock(spec=features.BoundingBox),
+            # labels, bboxes, masks
+            mocker.MagicMock(spec=features.BoundingBox),
+        ]
+
+        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", [features.Image, PIL.Image.Image, torch.Tensor])
+    def test__extract_image_targets(self, image_type, mocker):
+        transform = transforms.MixupDetection()
+
+        flat_sample = [
+            # images, batch size = 2
+            self.create_fake_image(mocker, image_type),
+            self.create_fake_image(mocker, image_type),
+            # labels, bboxes
+            mocker.MagicMock(spec=features.Label),
+            mocker.MagicMock(spec=features.BoundingBox),
+            # labels, bboxes
+            mocker.MagicMock(spec=features.Label),
+            mocker.MagicMock(spec=features.BoundingBox),
+        ]
+
+        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], to_tensor(flat_sample[0]))
+            torch.testing.assert_close(images[1], to_tensor(flat_sample[1]))
+        else:
+            assert images[0] == flat_sample[0]
+            assert images[1] == flat_sample[1]
+
+    def test__mixup(self):
+        image1 = 2 * torch.ones(3, 32, 64)
+        target_1 = {
+            "boxes": features.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": features.Label(torch.tensor([1, 2])),
+        }
+
+        image2 = 10 * torch.ones(3, 64, 32)
+        target_2 = {
+            "boxes": features.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": features.Label(torch.tensor([2, 3])),
+        }
+
+        transform = transforms.MixupDetection()
+        output_image, output_target = transform._mixup(image1, target_1, image2, target_2)
+        assert output_image.shape == (3, 64, 64)
+        assert output_target["boxes"].spatial_size == (64, 64)
+        assert len(output_target["boxes"]) == 4
+        assert len(output_target["labels"]) == 4

torchvision/prototype/datasets/_builtin/coco.py

@@ -24,7 +24,7 @@
     path_accessor,
     read_categories_file,
 )
-from torchvision.prototype.features import _Feature, BoundingBox, EncodedImage, Label
+from torchvision.prototype.features import _Feature, BoundingBox, EncodedImage, Label, Mask
 
 from .._api import register_dataset, register_info
 
@@ -114,7 +114,7 @@ def _decode_instances_anns(self, anns: List[Dict[str, Any]], image_meta: Dict[st
         labels = [ann["category_id"] for ann in anns]
         return dict(
             # TODO: create a segmentation feature
-            segmentations=_Feature(
+            segmentations=Mask(
                 torch.stack(
                     [
                         self._segmentation_to_mask(

torchvision/prototype/transforms/__init__.py

@@ -51,6 +51,7 @@
     ToDtype,
     TransposeDimensions,
 )
+from ._mixup_detection import MixupDetection
 from ._temporal import UniformTemporalSubsample
 from ._type_conversion import DecodeImage, LabelToOneHot, PILToTensor, ToImagePIL, ToImageTensor, ToPILImage
 

torchvision/prototype/transforms/_augment.py

@@ -1,17 +1,17 @@
 import math
 import numbers
 import warnings
-from typing import Any, cast, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, cast, Dict, List, Optional, Tuple, Type, Union
 
 import PIL.Image
 import torch
-from torch.utils._pytree import tree_flatten, tree_unflatten
+from torch.utils._pytree import tree_flatten, tree_unflatten, TreeSpec
 from torchvision.ops import masks_to_boxes
 from torchvision.prototype import features
 from torchvision.prototype.transforms import functional as F, InterpolationMode
 
 from ._transform import _RandomApplyTransform
-from ._utils import has_any, query_chw, query_spatial_size
+from ._utils import _isinstance, has_any, query_chw, query_spatial_size
 
 
 class RandomErasing(_RandomApplyTransform):
@@ -190,6 +190,90 @@ def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
             return inpt
 
 
+def flatten_and_extract_data(
+    inputs: Any, **target_types_or_checks: Tuple[Union[Type, Callable[[Any], bool]], ...]
+) -> Tuple[Tuple[List[Any], TreeSpec, List[Dict[str, int]]], List[features.TensorImageType], List[Dict[str, Any]]]:
+    # Images are special in the sense that they will always be extracted and returned
+    # separately. Internally however, they behave just as the other features.
+    types_or_checks: Dict[str, Tuple[Union[Type, Callable[[Any], bool]], ...]] = {
+        "images": (features.Image, PIL.Image.Image, features.is_simple_tensor),
+        **target_types_or_checks,
+    }
+
+    batch = inputs if len(inputs) > 1 else inputs[0]
+    flat_batch = []
+    sample_specs = []
+
+    offset = 0
+    batch_idcs = []
+    batch_data = []
+    for sample_idx, sample in enumerate(batch):
+        flat_sample, sample_spec = tree_flatten(sample)
+        flat_batch.extend(flat_sample)
+        sample_specs.append(sample_spec)
+
+        sample_types_or_checks = types_or_checks.copy()
+        sample_idcs = {}
+        sample_data = {}
+        for flat_idx, item in enumerate(flat_sample, offset):
+            if not sample_types_or_checks:
+                break
+
+            for key, types_or_checks_ in sample_types_or_checks.items():
+                if _isinstance(item, types_or_checks_):
+                    break
+            else:
+                continue
+
+            del sample_types_or_checks[key]
+            sample_idcs[key] = flat_idx
+            sample_data[key] = item
+
+        if sample_types_or_checks:
+            # TODO: improve message
+            raise TypeError(f"Sample at index {sample_idx} in the batch is missing {sample_types_or_checks.keys()}`")
+
+        batch_idcs.append(sample_idcs)
+        batch_data.append(sample_data)
+        offset += len(flat_sample)
+
+    batch_spec = TreeSpec(list, context=None, children_specs=sample_specs)
+
+    targets = batch_data
+    batch_data = []
+    for target in targets:
+        image = target.pop("images")
+        if isinstance(image, features.Image):
+            image = image.as_subclass(torch.Tensor)
+        elif isinstance(image, PIL.Image.Image):
+            image = F.pil_to_tensor(image)
+        batch_data.append(image)
+
+    return (flat_batch, batch_spec, batch_idcs), batch_data, targets
+
+
+def unflatten_and_insert_data(
+    flat_batch_with_spec: Tuple[List[Any], TreeSpec, List[Dict[str, int]]],
+    images: List[features.TensorImageType],
+    targets: List[Dict[str, Any]],
+) -> Any:
+    flat_batch, batch_spec, batch_idcs = flat_batch_with_spec
+
+    for sample_idx, sample_idcs in enumerate(batch_idcs):
+        for key, flat_idx in sample_idcs.items():
+            item = images[sample_idx] if key == "images" else targets[sample_idx][key]
+
+            inpt = flat_batch[flat_idx]
+            if isinstance(inpt, features._Feature):
+                item = type(inpt).wrap_like(inpt, item)
+            elif isinstance(inpt, PIL.Image.Image):
+                item = F.to_image_pil(item)
+
+            flat_batch[flat_idx] = item
+
+    return tree_unflatten(flat_batch, batch_spec)
+
+
 class SimpleCopyPaste(_RandomApplyTransform):
     def __init__(
         self,
@@ -214,7 +298,6 @@ def _copy_paste(
         resize_interpolation: F.InterpolationMode,
         antialias: Optional[bool],
     ) -> Tuple[features.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 +324,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 +364,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[features.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, features.Image) or features.is_simple_tensor(obj):
-                images.append(obj)
-            elif isinstance(obj, PIL.Image.Image):
-                images.append(F.to_image_tensor(obj))
-            elif isinstance(obj, features.BoundingBox):
-                bboxes.append(obj)
-            elif isinstance(obj, features.Mask):
-                masks.append(obj)
-            elif isinstance(obj, (features.Label, features.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[features.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, features.Image):
-                flat_sample[i] = features.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 features.is_simple_tensor(obj):
-                flat_sample[i] = output_images[c0]
-                c0 += 1
-            elif isinstance(obj, features.BoundingBox):
-                flat_sample[i] = features.BoundingBox.wrap_like(obj, output_targets[c1]["boxes"])
-                c1 += 1
-            elif isinstance(obj, features.Mask):
-                flat_sample[i] = features.Mask.wrap_like(obj, output_targets[c2]["masks"])
-                c2 += 1
-            elif isinstance(obj, (features.Label, features.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=(features.BoundingBox,),
+            masks=(features.Mask,),
+            labels=(features.Label, features.OneHotLabel),
+        )
 
         # images = [t1, t2, ..., tN]
         # Let's define paste_images as shifted list of input images
@@ -380,7 +409,4 @@ 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 tree_unflatten(flat_inputs, spec)
+        return unflatten_and_insert_data(flat_batch_with_spec, output_images, output_targets)