Add SetClampingMode transform

Author: NicolasHug

test/common_utils.py

@@ -410,7 +410,7 @@ def make_bounding_boxes(
     canvas_size=DEFAULT_SIZE,
     *,
     format=tv_tensors.BoundingBoxFormat.XYXY,
-    clamping_mode="soft",
+    clamping_mode="hard",  # TODOBB
     num_boxes=1,
     dtype=None,
     device="cpu",
@@ -481,7 +481,7 @@ def sample_position(values, max_value):
             out_boxes[:, :2] += buffer // 2
         elif format is tv_tensors.BoundingBoxFormat.XYXYXYXY:
             out_boxes[:, :] += buffer // 2
-    return tv_tensors.BoundingBoxes(out_boxes, format=format, canvas_size=canvas_size)
+    return tv_tensors.BoundingBoxes(out_boxes, format=format, canvas_size=canvas_size, clamping_mode=clamping_mode)
 
 
 def make_detection_masks(size=DEFAULT_SIZE, *, num_masks=1, dtype=None, device="cpu"):

test/test_transforms_v2.py

@@ -5589,6 +5589,54 @@ def test_clamping_mode(self, rotated, constructor_clamping_mode, clamping_mode,
         else:
             assert_equal(out, expected_clamped_output)
 
+class TestSetClampingMode:
+
+    @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))
+    @pytest.mark.parametrize("constructor_clamping_mode", ("hard", "none"))  # TODOBB add soft
+    @pytest.mark.parametrize("desired_clamping_mode", ("hard", "none"))  # TODOBB add soft
+    def test_setter(self, format, constructor_clamping_mode, desired_clamping_mode):
+
+        in_boxes = make_bounding_boxes(format=format, clamping_mode=constructor_clamping_mode)
+        out_boxes = transforms.SetClampingMode(clamping_mode=desired_clamping_mode)(in_boxes)
+
+        assert in_boxes.clamping_mode == constructor_clamping_mode  # input is unchanged: no leak
+        assert out_boxes.clamping_mode == desired_clamping_mode
+    
+    @pytest.mark.parametrize("format", list(tv_tensors.BoundingBoxFormat))
+    @pytest.mark.parametrize("constructor_clamping_mode", ("hard", "none"))  # TODOBB add soft
+    def test_pipeline_no_leak(self, format, constructor_clamping_mode):
+
+        class AssertClampingMode(transforms.Transform):
+            def __init__(self, expected_clamping_mode):
+                super().__init__()
+                self.expected_clamping_mode = expected_clamping_mode
+
+            _transformed_types = (tv_tensors.BoundingBoxes,)
+            
+            def transform(self, inpt, _):
+                assert inpt.clamping_mode == self.expected_clamping_mode
+                return inpt
+
+        t = transforms.Compose(
+            [
+                transforms.SetClampingMode("none"),
+                AssertClampingMode("none"),
+                transforms.SetClampingMode("hard"),
+                AssertClampingMode("hard"),
+                transforms.SetClampingMode("none"),
+                AssertClampingMode("none"),
+                transforms.ClampBoundingBoxes("hard")
+            ]
+        )
+
+        in_boxes = make_bounding_boxes(format=format, clamping_mode=constructor_clamping_mode)
+        out_boxes = t(in_boxes)
+
+        assert in_boxes.clamping_mode == constructor_clamping_mode  # input is unchanged: no leak
+
+        # assert that the output boxes clamping_mode is the one set by the last SetClampingMode.
+        # ClampBoundingBoxes doesn't set clamping_mode.
+        assert out_boxes.clamping_mode == "none"
 
 
 class TestClampKeyPoints:
@@ -7376,3 +7424,4 @@ def test_different_sizes(self, make_input1, make_input2, query):
     def test_no_valid_input(self, query):
         with pytest.raises(TypeError, match="No image"):
             query(["blah"])
+

torchvision/transforms/v2/__init__.py

@@ -41,7 +41,7 @@
     ScaleJitter,
     TenCrop,
 )
-from ._meta import ClampBoundingBoxes, ClampKeyPoints, ConvertBoundingBoxFormat
+from ._meta import ClampBoundingBoxes, ClampKeyPoints, ConvertBoundingBoxFormat, SetClampingMode
 from ._misc import (
     ConvertImageDtype,
     GaussianBlur,

torchvision/transforms/v2/_meta.py

@@ -1,4 +1,4 @@
-from typing import Any, Union
+from typing import Any, Union, Optional
 
 from torchvision import tv_tensors
 from torchvision.transforms.v2 import functional as F, Transform
@@ -30,10 +30,10 @@ class ClampBoundingBoxes(Transform):
     The clamping is done according to the bounding boxes' ``canvas_size`` meta-data.
 
     Args:
-        clamping_mode: TODOBB more docs. Default is None which relies on the input box' .clamping_mode attribute.
+        clamping_mode: TODOBB more docs. Default is None which relies on the input box' clamping_mode attribute.
 
     """
-    def __init__(self, clamping_mode: CLAMPING_MODE_TYPE = None) -> None:
+    def __init__(self, clamping_mode: Optional[CLAMPING_MODE_TYPE] = None) -> None:
         super().__init__()
         self.clamping_mode = clamping_mode
 
@@ -53,3 +53,18 @@ class ClampKeyPoints(Transform):
 
     def transform(self, inpt: tv_tensors.KeyPoints, params: dict[str, Any]) -> tv_tensors.KeyPoints:
         return F.clamp_keypoints(inpt)  # type: ignore[return-value]
+
+
+class SetClampingMode(Transform):
+    """TODOBB"""
+    def __init__(self, clamping_mode: CLAMPING_MODE_TYPE) -> None:
+        super().__init__()
+        # TODOBB validate mode
+        self.clamping_mode = clamping_mode
+
+    _transformed_types = (tv_tensors.BoundingBoxes,)
+
+    def transform(self, inpt: tv_tensors.BoundingBoxes, params: dict[str, Any]) -> tv_tensors.BoundingBoxes:
+        out = inpt.clone()
+        out.clamping_mode = self.clamping_mode
+        return out

torchvision/tv_tensors/__init__.py

@@ -23,14 +23,15 @@ def wrap(wrappee, *, like, **kwargs):
         wrappee (Tensor): The tensor to convert.
         like (:class:`~torchvision.tv_tensors.TVTensor`): The reference.
             ``wrappee`` will be converted into the same subclass as ``like``.
-        kwargs: Can contain "format" and "canvas_size" if ``like`` is a :class:`~torchvision.tv_tensor.BoundingBoxes`.
+        kwargs: Can contain "format", "canvas_size" and "clamping_mode" if ``like`` is a :class:`~torchvision.tv_tensor.BoundingBoxes`.
             Ignored otherwise.
     """
     if isinstance(like, BoundingBoxes):
         return BoundingBoxes._wrap(
             wrappee,
             format=kwargs.get("format", like.format),
             canvas_size=kwargs.get("canvas_size", like.canvas_size),
+            clamping_mode=kwargs.get("clamping_mode", like.clamping_mode),
         )
     elif isinstance(like, KeyPoints):
         return KeyPoints._wrap(wrappee, canvas_size=kwargs.get("canvas_size", like.canvas_size))