From a081ade167af0f64a17395cc3b18de4a3e44d892 Mon Sep 17 00:00:00 2001 From: YeonwooSung Date: Sat, 5 Sep 2026 22:49:53 +0900 Subject: [PATCH] Allow non-tensor labels in SanitizeBoundingBoxes and SanitizeKeyPoints --- test/test_transforms_v2.py | 86 +++++++++++++++++++-- torchvision/transforms/v2/_misc.py | 115 +++++++++++++++++++---------- 2 files changed, 154 insertions(+), 47 deletions(-) diff --git a/test/test_transforms_v2.py b/test/test_transforms_v2.py index 6d9f9732552..c55032e6d9b 100644 --- a/test/test_transforms_v2.py +++ b/test/test_transforms_v2.py @@ -7554,14 +7554,62 @@ def test_errors_transform(self): bad_labels_key = {"bbox": good_bbox, "BAD_KEY": torch.arange(good_bbox.shape[0])} transforms.SanitizeBoundingBoxes()(bad_labels_key) - with pytest.raises(ValueError, match="must be a tensor"): - not_a_tensor = {"bbox": good_bbox, "labels": torch.arange(good_bbox.shape[0]).tolist()} - transforms.SanitizeBoundingBoxes()(not_a_tensor) + with pytest.raises(ValueError, match="must be a tensor, numpy array, sequence"): + transforms.SanitizeBoundingBoxes()({"bbox": good_bbox, "labels": 0}) + + with pytest.raises(ValueError, match="must be a tensor, numpy array, sequence"): + transforms.SanitizeBoundingBoxes(labels_getter=lambda sample: (sample["labels"], None))( + {"bbox": good_bbox, "labels": torch.arange(good_bbox.shape[0])} + ) with pytest.raises(ValueError, match="Number of boxes"): different_sizes = {"bbox": good_bbox, "labels": torch.arange(good_bbox.shape[0] + 3)} transforms.SanitizeBoundingBoxes()(different_sizes) + @pytest.mark.parametrize( + "make_labels", + ( + lambda n: torch.arange(n), + lambda n: torch.arange(n).tolist(), + lambda n: tuple(range(n)), + lambda n: np.arange(n), + ), + ids=("tensor", "list", "tuple", "numpy"), + ) + def test_indexable_labels(self, make_labels): + H, W, min_size, min_area = 256, 128, 10, 10 + boxes, expected_valid_mask = self._get_boxes_and_valid_mask(H=H, W=W, min_size=min_size, min_area=min_area) + valid_indices = [i for (i, is_valid) in enumerate(expected_valid_mask) if is_valid] + labels = make_labels(boxes.shape[0]) + + out = transforms.SanitizeBoundingBoxes(min_size=min_size, min_area=min_area)( + {"boxes": boxes, "labels": labels} + ) + + assert type(out["labels"]) is type(labels) + assert list(out["labels"]) == valid_indices + assert out["boxes"].shape[0] == len(valid_indices) + + def test_multiple_indexable_label_entries(self): + H, W, min_size, min_area = 256, 128, 10, 10 + boxes, expected_valid_mask = self._get_boxes_and_valid_mask(H=H, W=W, min_size=min_size, min_area=min_area) + valid_indices = [i for (i, is_valid) in enumerate(expected_valid_mask) if is_valid] + sample = { + "boxes": boxes, + "labels": list(range(boxes.shape[0])), + "other_labels": np.arange(boxes.shape[0]), + } + + out = transforms.SanitizeBoundingBoxes( + min_size=min_size, + min_area=min_area, + labels_getter=lambda inputs: (inputs["labels"], inputs["other_labels"]), + )(sample) + + assert out["labels"] == valid_indices + assert isinstance(out["other_labels"], np.ndarray) + assert list(out["other_labels"]) == valid_indices + def test_errors_functional(self): good_bbox = tv_tensors.BoundingBoxes( @@ -7843,16 +7891,42 @@ def test_errors_transform(self): bad_sample = {"keypoints": good_keypoints, "BAD_KEY": torch.tensor([0])} transforms.SanitizeKeyPoints(labels_getter="default")(bad_sample) - # Test labels not a tensor - with pytest.raises(ValueError, match="must be a tensor"): - bad_sample = {"keypoints": good_keypoints, "labels": [0]} + with pytest.raises(ValueError, match="must be a tensor, numpy array, sequence"): + bad_sample = {"keypoints": good_keypoints, "labels": 0} transforms.SanitizeKeyPoints(labels_getter="default")(bad_sample) + with pytest.raises(ValueError, match="must be a tensor, numpy array, sequence"): + bad_sample = {"keypoints": good_keypoints, "labels": torch.tensor([0])} + transforms.SanitizeKeyPoints(labels_getter=lambda sample: (sample["labels"], None))(bad_sample) + # Test mismatched sizes with pytest.raises(ValueError, match="Number of"): bad_sample = {"keypoints": good_keypoints, "labels": torch.tensor([0, 1, 2])} transforms.SanitizeKeyPoints(labels_getter="default")(bad_sample) + @pytest.mark.parametrize( + "make_labels", + ( + lambda n: torch.arange(n), + lambda n: torch.arange(n).tolist(), + lambda n: tuple(range(n)), + lambda n: np.arange(n), + ), + ids=("tensor", "list", "tuple", "numpy"), + ) + def test_indexable_labels(self, make_labels): + canvas_size = (40, 40) + keypoints, expected_validity = self._make_keypoints_with_validity(canvas_size=canvas_size, shape="2d") + keypoints = tv_tensors.KeyPoints(keypoints, canvas_size=canvas_size) + valid_indices = [i for i, is_valid in enumerate(expected_validity) if is_valid] + labels = make_labels(keypoints.shape[0]) + + out = transforms.SanitizeKeyPoints(labels_getter="default")({"keypoints": keypoints, "labels": labels}) + + assert type(out["labels"]) is type(labels) + assert list(out["labels"]) == valid_indices + assert out["keypoints"].shape[0] == len(valid_indices) + def test_no_label(self): """Test transform without labels.""" img = make_image() diff --git a/torchvision/transforms/v2/_misc.py b/torchvision/transforms/v2/_misc.py index 305149c87b1..166c0ff5f8e 100644 --- a/torchvision/transforms/v2/_misc.py +++ b/torchvision/transforms/v2/_misc.py @@ -2,6 +2,7 @@ from collections.abc import Sequence from typing import Any, Callable, Optional, Union +import numpy as np import PIL.Image import torch @@ -20,6 +21,53 @@ is_pure_tensor, ) +_LABELS_TYPE_MSG = ( + "The labels in the input to forward() must be a tensor, numpy array, sequence, or None, got {type} instead." +) + + +def _get_entry_length(entry: Any) -> int: + shape = getattr(entry, "shape", None) + if shape is not None and len(shape) > 0: + return shape[0] + return len(entry) + + +def _index_with_valid(entry: Any, valid: torch.Tensor) -> Any: + if isinstance(entry, (list, tuple)): + return type(entry)(item for item, keep in zip(entry, valid.tolist()) if keep) + return entry[valid] + + +def _is_indexable_label_entry(entry: Any) -> bool: + return isinstance(entry, (torch.Tensor, np.ndarray, list, tuple)) + + +def _normalize_sanitize_labels(labels: Any) -> tuple[Any, ...]: + if isinstance(labels, (list, tuple)): + # A sequence of tensors/arrays/sequences is multiple label entries. A sequence of + # scalars (e.g. list[int]) is a single per-box/per-keypoint label entry. + if any(_is_indexable_label_entry(entry) for entry in labels): + for entry in labels: + if not _is_indexable_label_entry(entry): + raise ValueError(_LABELS_TYPE_MSG.format(type=type(entry))) + return tuple(labels) + return (labels,) + if _is_indexable_label_entry(labels): + return (labels,) + raise ValueError(_LABELS_TYPE_MSG.format(type=type(labels))) + + +def _labels_tree_is_leaf(labels: Optional[tuple[Any, ...]]) -> Optional[Callable[[Any], bool]]: + if labels is None: + return None + label_ids = {id(label) for label in labels} + + def is_leaf(x: Any) -> bool: + return id(x) in label_ids + + return is_leaf + # TODO: do we want/need to expose this? class Identity(Transform): @@ -393,8 +441,9 @@ class SanitizeBoundingBoxes(Transform): 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. + - A single tensor, numpy array, or sequence (the labels) + - A tuple/list of tensors, numpy arrays, or sequences, 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. @@ -425,26 +474,18 @@ def forward(self, *inputs: Any) -> Any: labels = self._labels_getter(inputs) 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)): - 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) + labels = _normalize_sanitize_labels(labels) + + # Keep list/tuple labels as leaves so they can be subset with the validity mask. + flat_inputs, spec = tree_flatten(inputs, is_leaf=_labels_tree_is_leaf(labels)) boxes = get_bounding_boxes(flat_inputs) if labels is not None: for label in labels: - if boxes.shape[0] != label.shape[0]: + if boxes.shape[0] != _get_entry_length(label): raise ValueError( f"Number of boxes (shape={boxes.shape}) and must match the number of labels." - f"Found labels with shape={label.shape})." + f"Found labels with length={_get_entry_length(label)}." ) valid = F._misc._get_sanitize_bounding_boxes_mask( @@ -468,16 +509,16 @@ def transform(self, inpt: Any, params: dict[str, Any]) -> Any: if not (is_label or is_bounding_boxes or is_mask): return inpt + if is_label: + return _index_with_valid(inpt, params["valid"]) + try: output = inpt[params["valid"]] except (IndexError): # If indexing fails (e.g., shape mismatch), pass through unchanged return inpt - if is_label: - return output - else: - return tv_tensors.wrap(output, like=inpt) + return tv_tensors.wrap(output, like=inpt) class SanitizeKeyPoints(Transform): @@ -504,8 +545,9 @@ class SanitizeKeyPoints(Transform): 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 keypoints. + - A single tensor, numpy array, or sequence (the labels) + - A tuple/list of tensors, numpy arrays, or sequences, each of which will be subject to the same + sanitization as the keypoints. If ``labels_getter`` is None (the default), then only keypoints are sanitized. """ @@ -523,26 +565,18 @@ def forward(self, *inputs: Any) -> Any: labels = self._labels_getter(inputs) 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)): - 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) + labels = _normalize_sanitize_labels(labels) + + # Keep list/tuple labels as leaves so they can be subset with the validity mask. + flat_inputs, spec = tree_flatten(inputs, is_leaf=_labels_tree_is_leaf(labels)) points = get_keypoints(flat_inputs) if labels is not None: for label in labels: - if points.shape[0] != label.shape[0]: + if points.shape[0] != _get_entry_length(label): raise ValueError( f"Number of kepyoints (shape={points.shape}) must match the number of labels." - f"Found labels with shape={label.shape})." + f"Found labels with length={_get_entry_length(label)}." ) valid = F._misc._get_sanitize_keypoints_mask( @@ -562,9 +596,8 @@ def transform(self, inpt: Any, params: dict[str, Any]) -> Any: if not (is_label or is_keypoints): return inpt - output = inpt[params["valid"]] - if is_label: - return output - else: - return tv_tensors.wrap(output, like=inpt) + return _index_with_valid(inpt, params["valid"]) + + output = inpt[params["valid"]] + return tv_tensors.wrap(output, like=inpt)