Repository: pytorch/vision
Pull request: Allow SanitizeBoundingBoxes to sanitize more labels
Closes https://github.com/pytorch/vision/issues/8294

We allow the `labels_getter` parameter of `SanitizeBoudingBoxes` to return a tuple. Each entry in the tuple gets sanitized/subsampled, along with the boxes. This enables the sanitization of properties like "area" or "iscrowd" from COCO.

cc @vfdev-5

Problem statement:
Cleaner SanitizeBoundingBoxes transform
### 🚀 The feature

- A functional version of SanitizeBoundingBoxes
- A version of SanitizeBoundingBoxes that is aligned with other detection transforms: it should either take a BoundingBoxes TVTensor, or other transforms should be able to take in a dictionary that contains boxes, among others.
- A version of SanitizeBoundingBoxes that can deal with more keys from that dictionary that just "boxes" and "labels"

### Motivation, pitch

SanitizeBoundingBoxes is a transform for bounding boxes that is a bit different from the others, as it doesn't take a bounding box object but a dictionary containing a bounding box object and maybe labels.

So we cannot use it in a `transforms.Compose` object as we can with other transforms, and have to apply it separately.
Plus, it has no functional version, which forces us to instantiate a transforms object and apply it to said dictionary.

Finally, an important drawback of SanitizeBoundingBoxes is that it subsets the boxes and the labels together, but not other detection features such as `area` or `iscrowd`. This is error prone, and this transform should at least be able to accommodate the standard COCO format.


### Alternatives

_No response_

### Additional context

_No response_

Hints from issue discussion:
Thanks a lot for the detailed feature request @MarcSzafraniec . This sounds reasonable overall, I shared more detailed thoughts below. LMK what you think!

> A functional version of SanitizeBoundingBoxes

Sounds good. Not sure whether we should allow `labels` and other keys here, or just bouding boxes, we'll see. Maybe we can get away by just returning the [ `valid` mask](https://github.com/pytorch/vision/blob/main/torchvision/transforms/v2/_misc.py#L385) and let the user subsample they input.


> A version of SanitizeBoundingBoxes that is aligned with other detection transforms: it should either take a BoundingBoxes TVTensor, or other transforms should be able to take in a dictionary that contains boxes, among others.

All of this is already supported. I.e. `SanitizeBoundingBoxes()` supports a single `BoundingBoxes()` object as input:

```py
from torchvision.transforms import v2
from torchvision import tv_tensors
import torch

bbox = tv_tensors.BoundingBoxes(torch.tensor([[0, 0, 10, 10], [10, 10, 20, 20]]), canvas_size=(100, 100), format="XYXY")
v2.SanitizeBoundingBoxes(labels_getter=None)(bbox)

# BoundingBoxes([[ 0,  0, 10, 10],
#                [10, 10, 20, 20]], format=BoundingBoxFormat.XYXY, canvas_size=(100, 100))
```

and the other transforms do support dictionary and arbitrarily nested inputs, just like `SanitizeBoundingBoxes`: https://pytorch.org/vision/main/auto_examples/transforms/plot_transforms_getting_started.html#what-do-i-pass-as-input

You can definitely use `SanitizeBoundingBoxes()` within ` Compose()` for example: https://pytorch.org/vision/main/auto_examples/transforms/plot_transforms_e2e.html#transforms

It seems like something wasn't clear in the docs here - I'd love for you to let us know what could be improved on that front.


> Finally, an important drawback of SanitizeBoundingBoxes is that it subsets the boxes and the labels together, but not other detection features such as area or iscrowd. This is error prone, and this transform should at least be able to accommodate the standard COCO format.

In general it's going to be untracktable to subset arbitrary inputs (and potentially inconvenient, since it would force those inputs to be in a specific expected format). But perhaps it would be enough let the functional return the subsampling maks (as suggested above), and let the user do the indexing themselves on their desired ipnuts (area, iscrowd, etc.) ?
> Thanks a lot for the detailed feature request @MarcSzafraniec . This sounds reasonable overall, I shared more detailed thoughts below. LMK what you think!
> 
> > A functional version of SanitizeBoundingBoxes
> 
> Sounds good. Not sure whether we should allow `labels` and other keys here, or just bouding boxes, we'll see. Maybe we can get away by just returning the [ `valid` mask](https://github.com/pytorch/vision/blob/main/torchvision/transforms/v2/_misc.py#L385) and let the user subsample they input.
> 
> > A version of SanitizeBoundingBoxes that is aligned with other detection transforms: it should either take a BoundingBoxes TVTensor, or other transforms should be able to take in a dictionary that contains boxes, among others.
> 
> All of this is already supported. I.e. `SanitizeBoundingBoxes()` supports a single `BoundingBoxes()` object as input:
> 
> ```python
> from torchvision.transforms import v2
> from torchvision import tv_tensors
> import torch
> 
> bbox = tv_tensors.BoundingBoxes(torch.tensor([[0, 0, 10, 10], [10, 10, 20, 20]]), canvas_size=(100, 100), format="XYXY")
> v2.SanitizeBoundingBoxes(labels_getter=None)(bbox)
> 
> # BoundingBoxes([[ 0,  0, 10, 10],
> #                [10, 10, 20, 20]], format=BoundingBoxFormat.XYXY, canvas_size=(100, 100))
> ```
> 
> and the other transforms do support dictionary and arbitrarily nested inputs, just like `SanitizeBoundingBoxes`: https://pytorch.org/vision/main/auto_examples/transforms/plot_transforms_getting_started.html#what-do-i-pass-as-input
> 
> You can definitely use `SanitizeBoundingBoxes()` within ` Compose()` for example: https://pytorch.org/vision/main/auto_examples/transforms/plot_transforms_e2e.html#transforms
> 
> It seems like something wasn't clear in the docs here - I'd love for you to let us know what could be improved on that front.
> 
> > Finally, an important drawback of SanitizeBoundingBoxes is that it subsets the boxes and the labels together, but not other detection features such as area or iscrowd. This is error prone, and this transform should at least be able to accommodate the standard COCO format.
> 
> In general it's going to be untracktable to subset arbitrary inputs (and potentially inconvenient, since it would force those inputs to be in a specific expected format). But perhaps it would be enough let the functional return the subsampling maks (as suggested above), and let the user do the indexing themselves on their desired ipnuts (area, iscrowd, etc.) ?

Thanks for answering so quickly ! Well, indeed we can pass either a `BoundingBoxes` object to `SanitizeBoundingBoxes`, or a dict to other transforms, but then the `mask` information is lost. 

Maybe a solution would be to have a "labels_getter" that could be in fact a list of dictionary keys for which to apply the mask ? It would be simple to implement and probably generic enough.
Summarizing an offline chat with @MarcSzafraniec : 


### A. Allow users to subsample arbitrary stuff

We need to allow users to subsample arbitrary stuff, not just the labels (e.g. "iscrowd" info, etc). This can be done in different ways, with the constraint that we don't want to enforce a specific input structure as this would go against a fundamental design principle of v2 (e.g. we don't want to enforce specific keys or even a dict structure).

One suggested solution would be to let `labels_getter` return a tuple of tensors. Each tensor in that tuple gets subsampled, just like the label tensor. The tuple entries technically don't need to be tensors, they just have to be anything that supports `[bool_mask]` indexing.

### B. Having a functional

We've left out the functional implementation of `SanitizeBoundingBoxes`, but it can still be useful. This functional should be able to subsample the `BoundingBoxes` objects but also the labels and arbitrary tensors. I see 2 main ways of doing this for now (there's probably others), both of which deviate from what we typically call a `functional` (then again, this `SanitizeBoundingBoxes` is an outlier itself)

Option 1

```py
def sanitize_bounding_boxes(boxes, *args):
    # subsample boxes and everything in *args

# User code:
boxes, labels, iscrowd = sanitize_bounding_boxes(boxes, labels, is_crowd)
```

Option 2

```py
def get_sanitize_bouding_boxes_mask(boxes):
    # just return the boolean subsampling mask
	return is_valid

# User code:
valid_mask = get_sanitize_bouding_boxes_mask(boxes)
boxes, labels, iscrowd = boxes[valid_mask], labels[valid_mask], iscrowd[valid_mask]
```

Note: a way to get the `valid_mask` with option 1 is to pass `boxes, valid_indices = sanitize_bounding_boxes(boxes, torch.arange(boxes.shape[0]))`

EDIT: actually Option2 is not as pleasant as I thought because it forces users to call `boxes = tv_tensors.wrap(...)` to preserve the BoundingBoxes type.


CC @pmeier @vfdev-5 for your input on both design discussions here (A and B).
A: 

:+1:

---

B: 

I like option 1 best, but I don't think JIT supports `*args`. So this is another way this would deviate from our current implementation. 

Plus, when you say "functional" here, do you mean an actual functional, i.e. `TVTensor` input or a kernel? In the former case how would you handle rewrapping?

Without breaking any of our current paradigms, let me propose a third option:

```python
def sanitize_bounding_boxes(boxes):
    # compute subsampling mask and apply to boxes
    return boxes, is_valid
```

That way people can use the mask if needed without extra computation, and don't have to subsample their boxes themselves. If we treat this function as kernel rather than as "functional", we also don't have an issue with the tuple return. For boxes we already have kernels that return a tuple, with the first item always being the boxes and the remaining parts being metadata.
I have an additional comment: why this check `if min_size < 1:            raise ValueError(f"min_size must be >= 1, got {min_size}.")` ? Boxes could be relative, or even just floating point, in that case its possible to have a width or height < 1 ...
@MarcSzafraniec that's because in torchvision the expected format for the bounding boxes coordinates is absolute, not relative. The rest of the transforms also expect e.g. the XYXY format to be in absolute pixel coordinates.

Commit message: Merge branch 'main' of github.com:pytorch/vision into sanitize_labels_tuple

Patch to review:
diff --git a/test/test_transforms_v2.py b/test/test_transforms_v2.py
index 743fb5d86bf..da4ba3141d5 100644
--- a/test/test_transforms_v2.py
+++ b/test/test_transforms_v2.py
@@ -5706,7 +5706,16 @@ def _get_boxes_and_valid_mask(self, H=256, W=128, min_size=10):
         return boxes, expected_valid_mask
 
     @pytest.mark.parametrize("min_size", (1, 10))
-    @pytest.mark.parametrize("labels_getter", ("default", lambda inputs: inputs["labels"], None, lambda inputs: None))
+    @pytest.mark.parametrize(
+        "labels_getter",
+        (
+            "default",
+            lambda inputs: inputs["labels"],
+            lambda inputs: (inputs["labels"], inputs["other_labels"]),
+            None,
+            lambda inputs: None,
+        ),
+    )
     @pytest.mark.parametrize("sample_type", (tuple, dict))
     def test_transform(self, min_size, labels_getter, sample_type):
 
@@ -5721,12 +5730,16 @@ def test_transform(self, min_size, labels_getter, sample_type):
 
         labels = torch.arange(boxes.shape[0])
         masks = tv_tensors.Mask(torch.randint(0, 2, size=(boxes.shape[0], H, W)))
+        # other_labels corresponds to properties from COCO like iscrowd, area...
+        # We only sanitize it when labels_getter returns a tuple
+        other_labels = torch.arange(boxes.shape[0])
         whatever = torch.rand(10)
         input_img = torch.randint(0, 256, size=(1, 3, H, W), dtype=torch.uint8)
         sample = {
             "image": input_img,
             "labels": labels,
             "boxes": boxes,
+            "other_labels": other_labels,
             "whatever": whatever,
             "None": None,
             "masks": masks,
@@ -5741,12 +5754,14 @@ def test_transform(self, min_size, labels_getter, sample_type):
         if sample_type is tuple:
             out_image = out[0]
             out_labels = out[1]["labels"]
+            out_other_labels = out[1]["other_labels"]
             out_boxes = out[1]["boxes"]
             out_masks = out[1]["masks"]
             out_whatever = out[1]["whatever"]
         else:
             out_image = out["image"]
             out_labels = out["labels"]
+            out_other_labels = out["other_labels"]
             out_boxes = out["boxes"]
             out_masks = out["masks"]
             out_whatever = out["whatever"]
@@ -5757,14 +5772,20 @@ def test_transform(self, min_size, labels_getter, sample_type):
         assert isinstance(out_boxes, tv_tensors.BoundingBoxes)
         assert isinstance(out_masks, tv_tensors.Mask)
 
-        if labels_getter is None or (callable(labels_getter) and labels_getter({"labels": "blah"}) is None):
+        if labels_getter is None or (callable(labels_getter) and labels_getter(sample) is None):
             assert out_labels is labels
+            assert out_other_labels is other_labels
         else:
             assert isinstance(out_labels, torch.Tensor)
             assert out_boxes.shape[0] == out_labels.shape[0] == out_masks.shape[0]
             # This works because we conveniently set labels to arange(num_boxes)
             assert out_labels.tolist() == valid_indices
 
+            if callable(labels_getter) and type(labels_getter(sample)) is tuple:
+                assert_equal(out_other_labels, out_labels)
+            else:
+                assert_equal(out_other_labels, other_labels)
+
     @pytest.mark.parametrize("input_type", (torch.Tensor, tv_tensors.BoundingBoxes))
     def test_functional(self, input_type):
         # Note: the "functional" F.sanitize_bounding_boxes was added after the class, so there is some
diff --git a/torchvision/transforms/v2/_misc.py b/torchvision/transforms/v2/_misc.py
index 852acc555c6..00bad2be253 100644
--- a/torchvision/transforms/v2/_misc.py
+++ b/torchvision/transforms/v2/_misc.py
@@ -321,6 +321,9 @@ class SanitizeBoundingBoxes(Transform):
     - have any coordinate outside of their corresponding image. You may want to
       call :class:`~torchvision.transforms.v2.ClampBoundingBoxes` first to avoid undesired removals.
 
+    It can also sanitize other tensors like the "iscrowd" or "area" properties from COCO
+    (see ``labels_getter`` parameter).
+
     It is recommended to call it at the end of a pipeline, before passing the
     input to the models. It is critical to call this transform if
     :class:`~torchvision.transforms.v2.RandomIoUCrop` was called.
@@ -330,12 +333,18 @@ class SanitizeBoundingBoxes(Transform):
 
     Args:
         min_size (float, optional) The size below which bounding boxes are removed. Default is 1.
-        labels_getter (callable or str or None, optional): indicates how to identify the labels in the input.
+        labels_getter (callable or str or None, optional): indicates how to identify the labels in the input
+            (or anything else that needs to be sanitized along with the bounding boxes).
             By default, this will try to find a "labels" key in the input (case-insensitive), if
             the input is a dict or it is a tuple whose second element is a dict.
             This heuristic should work well with a lot of datasets, including the built-in torchvision datasets.
-            It can also be a callable that takes the same input
-            as the transform, and returns the labels.
+
+            It can also be a callable that takes the same input as the transform, and returns either:
+
+            - A single tensor (the labels)
+            - A tuple/list of tensors, each of which will be subject to the same sanitization as the bounding boxes.
+              This is useful to sanitize multiple tensors like the labels, and the "iscrowd" or "area" properties
+              from COCO.
     """
 
     def __init__(
@@ -356,18 +365,29 @@ def forward(self, *inputs: Any) -> Any:
         inputs = inputs if len(inputs) > 1 else inputs[0]
 
         labels = self._labels_getter(inputs)
-        if labels is not None and not isinstance(labels, torch.Tensor):
-            raise ValueError(
-                f"The labels in the input to forward() must be a tensor or None, got {type(labels)} instead."
-            )
+        if labels is not None:
+            msg = "The labels in the input to forward() must be a tensor or None, got {type} instead."
+            if isinstance(labels, torch.Tensor):
+                labels = (labels,)
+            elif isinstance(labels, (tuple, list)):
+                labels = tuple(labels)
+                for entry in labels:
+                    if not isinstance(entry, torch.Tensor):
+                        # TODO: we don't need to enforce tensors, just that entries are indexable as t[bool_mask]
+                        raise ValueError(msg.format(type=type(entry)))
+            else:
+                raise ValueError(msg.format(type=type(labels)))
 
         flat_inputs, spec = tree_flatten(inputs)
         boxes = get_bounding_boxes(flat_inputs)
 
-        if labels is not None and boxes.shape[0] != labels.shape[0]:
-            raise ValueError(
-                f"Number of boxes (shape={boxes.shape}) and number of labels (shape={labels.shape}) do not match."
-            )
+        if labels is not None:
+            for label in labels:
+                if boxes.shape[0] != label.shape[0]:
+                    raise ValueError(
+                        f"Number of boxes (shape={boxes.shape}) and must match the number of labels."
+                        f"Found labels with shape={label.shape})."
+                    )
 
         valid = F._misc._get_sanitize_bounding_boxes_mask(
             boxes,
@@ -381,7 +401,7 @@ def forward(self, *inputs: Any) -> Any:
         return tree_unflatten(flat_outputs, spec)
 
     def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
-        is_label = inpt is not None and inpt is params["labels"]
+        is_label = params["labels"] is not None and any(inpt is label for label in params["labels"])
         is_bounding_boxes_or_mask = isinstance(inpt, (tv_tensors.BoundingBoxes, tv_tensors.Mask))
 
         if not (is_label or is_bounding_boxes_or_mask):
@@ -391,5 +411,5 @@ def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
 
         if is_label:
             return output
-
-        return tv_tensors.wrap(output, like=inpt)
+        else:
+            return tv_tensors.wrap(output, like=inpt)


Code review comments:
diff_hunk: @@ -356,18 +365,29 @@ def forward(self, *inputs: Any) -> Any:
         inputs = inputs if len(inputs) > 1 else inputs[0]
 
         labels = self._labels_getter(inputs)
-        if labels is not None and not isinstance(labels, torch.Tensor):
-            raise ValueError(
-                f"The labels in the input to forward() must be a tensor or None, got {type(labels)} instead."
-            )
+        if labels is not None:
+            msg = "The labels in the input to forward() must be a tensor or None, got {type} instead."
+            if isinstance(labels, torch.Tensor):
+                labels = (labels,)
+            elif isinstance(labels, (tuple, list)):
+                labels = tuple(labels)
path: torchvision/transforms/v2/_misc.py
text: @user1:
Why force into `tuple` here when we have already established that it can only be a tuple or list above?