fix tests

Author: AntoineSimoulin

test/common_utils.py

@@ -400,10 +400,10 @@ def make_image_pil(*args, **kwargs):
     return to_pil_image(make_image(*args, **kwargs))
 
 
-def make_keypoints(canvas_size=DEFAULT_SIZE, *, num_points=4, dtype=None, device="cpu"):
+def make_keypoints(canvas_size=DEFAULT_SIZE, *, clamping_mode="soft", num_points=4, dtype=None, device="cpu"):
     y = torch.randint(0, canvas_size[0], size=(num_points, 1), dtype=dtype, device=device)
     x = torch.randint(0, canvas_size[1], size=(num_points, 1), dtype=dtype, device=device)
-    return tv_tensors.KeyPoints(torch.cat((x, y), dim=-1), canvas_size=canvas_size)
+    return tv_tensors.KeyPoints(torch.cat((x, y), dim=-1), canvas_size=canvas_size, clamping_mode=clamping_mode)
 
 
 def make_bounding_boxes(

test/test_transforms_v2.py

@@ -633,6 +633,7 @@ def affine_rotated_bounding_boxes(bounding_boxes):
 
 def reference_affine_keypoints_helper(keypoints, *, affine_matrix, new_canvas_size=None, clamp=True):
     canvas_size = new_canvas_size or keypoints.canvas_size
+    clamping_mode = keypoints.clamping_mode
 
     def affine_keypoints(keypoints):
         dtype = keypoints.dtype
@@ -652,15 +653,15 @@ def affine_keypoints(keypoints):
         )
 
         if clamp:
-            output = F.clamp_keypoints(output, canvas_size=canvas_size)
+            output = F.clamp_keypoints(output, canvas_size=canvas_size, clamping_mode=clamping_mode)
         else:
             dtype = output.dtype
 
         return output.to(dtype=dtype, device=device)
 
     return tv_tensors.KeyPoints(
         torch.cat([affine_keypoints(k) for k in keypoints.reshape(-1, 2).unbind()], dim=0).reshape(keypoints.shape),
-        canvas_size=canvas_size,
+        canvas_size=canvas_size, clamping_mode=clamping_mode
     )
 
 
@@ -3309,7 +3310,6 @@ def test_functional(self, make_input):
             (F.elastic_image, tv_tensors.Image),
             (F.elastic_mask, tv_tensors.Mask),
             (F.elastic_video, tv_tensors.Video),
-            (F.elastic_keypoints, tv_tensors.KeyPoints),
         ],
     )
     def test_functional_signature(self, kernel, input_type):
@@ -5325,6 +5325,7 @@ def test_correctness_perspective_bounding_boxes(self, startpoints, endpoints, fo
 
     def _reference_perspective_keypoints(self, keypoints, *, startpoints, endpoints):
         canvas_size = keypoints.canvas_size
+        clamping_mode = keypoints.clamping_mode
         dtype = keypoints.dtype
         device = keypoints.device
 
@@ -5364,13 +5365,15 @@ def perspective_keypoints(keypoints):
             return F.clamp_keypoints(
                 output,
                 canvas_size=canvas_size,
+                clamping_mode=clamping_mode
             ).to(dtype=dtype, device=device)
 
         return tv_tensors.KeyPoints(
             torch.cat([perspective_keypoints(k) for k in keypoints.reshape(-1, 2).unbind()], dim=0).reshape(
                 keypoints.shape
             ),
             canvas_size=canvas_size,
+            clamping_mode=clamping_mode,
         )
 
     @pytest.mark.parametrize(("startpoints", "endpoints"), START_END_POINTS)
@@ -5733,32 +5736,80 @@ def test_error(self):
 
 
 class TestClampKeyPoints:
+    @pytest.mark.parametrize("clamping_mode", ("soft", "hard", None))
     @pytest.mark.parametrize("dtype", [torch.int64, torch.float32])
     @pytest.mark.parametrize("device", cpu_and_cuda())
-    def test_kernel(self, dtype, device):
-        keypoints = make_keypoints(dtype=dtype, device=device)
+    def test_kernel(self, clamping_mode, dtype, device):
+        keypoints = make_keypoints(dtype=dtype, device=device, clamping_mode=clamping_mode)
         check_kernel(
             F.clamp_keypoints,
             keypoints,
             canvas_size=keypoints.canvas_size,
+            clamping_mode=clamping_mode,
         )
 
-    def test_functional(self):
-        check_functional(F.clamp_keypoints, make_keypoints())
+    @pytest.mark.parametrize("clamping_mode", ("soft", "hard", None))
+    def test_functional(self, clamping_mode):
+        check_functional(F.clamp_keypoints, make_keypoints(clamping_mode=clamping_mode))
 
     def test_errors(self):
         input_tv_tensor = make_keypoints()
         input_pure_tensor = input_tv_tensor.as_subclass(torch.Tensor)
 
-        with pytest.raises(ValueError, match="`canvas_size` has to be passed"):
+        with pytest.raises(ValueError, match="`canvas_size` and `clamping_mode` have to be passed."):
             F.clamp_keypoints(input_pure_tensor, canvas_size=None)
 
         with pytest.raises(ValueError, match="`canvas_size` must not be passed"):
             F.clamp_keypoints(input_tv_tensor, canvas_size=input_tv_tensor.canvas_size)
+        with pytest.raises(ValueError, match="clamping_mode must be soft,"):
+            F.clamp_keypoints(input_tv_tensor, clamping_mode="bad")
+        with pytest.raises(ValueError, match="clamping_mode must be soft,"):
+            transforms.ClampKeyPoints(clamping_mode="bad")(input_tv_tensor)
 
     def test_transform(self):
         check_transform(transforms.ClampKeyPoints(), make_keypoints())
 
+    @pytest.mark.parametrize("constructor_clamping_mode", ("soft", "hard", None))
+    @pytest.mark.parametrize("clamping_mode", ("soft", "hard", None, "auto"))
+    @pytest.mark.parametrize("pass_pure_tensor", (True, False))
+    @pytest.mark.parametrize("fn", [F.clamp_keypoints, transform_cls_to_functional(transforms.ClampKeyPoints)])
+    def test_clamping_mode(self, constructor_clamping_mode, clamping_mode, pass_pure_tensor, fn):
+        # This test checks 2 things:
+        # - That passing clamping_mode=None to the clamp_keypointss
+        #   functional (or to the class) relies on the box's `.clamping_mode`
+        #   attribute
+        # - That clamping happens when it should, and only when it should, i.e.
+        #   when the clamping mode is not None. It doesn't validate the
+        #   numerical results, only that clamping happened. For that, we create
+        #   a keypoints with large coordinates (100) inside of a small 10x10 image.
+
+        if pass_pure_tensor and fn is not F.clamp_keypoints:
+            # Only the functional supports pure tensors, not the class
+            return
+        if pass_pure_tensor and clamping_mode == "auto":
+            # cannot leave clamping_mode="auto" when passing pure tensor
+            return
+
+        keypoints = tv_tensors.KeyPoints(
+            [[0, 100], [0, 100]],canvas_size=(10, 10), clamping_mode=constructor_clamping_mode
+        )
+        expected_clamped_output = torch.tensor([[0, 9], [0, 9]]) if clamping_mode == "hard" else torch.tensor([[0, 100], [0, 100]])
+
+        if pass_pure_tensor:
+            out = fn(
+                keypoints.as_subclass(torch.Tensor),
+                canvas_size=keypoints.canvas_size,
+                clamping_mode=clamping_mode,
+            )
+        else:
+            out = fn(keypoints, clamping_mode=clamping_mode)
+
+        clamping_mode_prevailing = constructor_clamping_mode if clamping_mode == "auto" else clamping_mode
+        if clamping_mode_prevailing is None:
+            assert_equal(keypoints, out)  # should be a pass-through
+        else:
+            assert_equal(out, expected_clamped_output)
+
 
 class TestInvert:
     @pytest.mark.parametrize("dtype", [torch.uint8, torch.int16, torch.float32])