Add ClampKeyPoints and corresponding test

Author: NicolasHug

docs/source/transforms.rst

@@ -408,6 +408,7 @@ Miscellaneous
     v2.Lambda
     v2.SanitizeBoundingBoxes
     v2.ClampBoundingBoxes
+    v2.ClampKeyPoints
     v2.UniformTemporalSubsample
     v2.JPEG
 
@@ -421,6 +422,7 @@ Functionals
     v2.functional.erase
     v2.functional.sanitize_bounding_boxes
     v2.functional.clamp_bounding_boxes
+    v2.functional.clamp_keypoints
     v2.functional.uniform_temporal_subsample
     v2.functional.jpeg
 

test/test_transforms_v2.py

@@ -5444,6 +5444,34 @@ def test_errors(self):
     def test_transform(self):
         check_transform(transforms.ClampBoundingBoxes(), make_bounding_boxes())
 
+class TestClampKeyPoints:
+    @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)
+        check_kernel(
+            F.clamp_keypoints,
+            keypoints,
+            canvas_size=keypoints.canvas_size,
+        )
+
+    def test_functional(self):
+        check_functional(F.clamp_keypoints, make_keypoints())
+
+    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"):
+            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)
+
+    def test_transform(self):
+        check_transform(transforms.ClampKeyPoints(), make_keypoints())
+
+
 
 class TestInvert:
     @pytest.mark.parametrize("dtype", [torch.uint8, torch.int16, torch.float32])

torchvision/transforms/v2/__init__.py

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

torchvision/transforms/v2/_meta.py

@@ -34,3 +34,15 @@ class ClampBoundingBoxes(Transform):
 
     def transform(self, inpt: tv_tensors.BoundingBoxes, params: dict[str, Any]) -> tv_tensors.BoundingBoxes:
         return F.clamp_bounding_boxes(inpt)  # type: ignore[return-value]
+
+class ClampKeyPoints(Transform):
+    """Clamp keypoints to their corresponding image dimensions.
+
+    The clamping is done according to the keypoints' ``canvas_size`` meta-data.
+
+    """
+
+    _transformed_types = (tv_tensors.KeyPoints,)
+
+    def transform(self, inpt: tv_tensors.KeyPoints, params: dict[str, Any]) -> tv_tensors.KeyPoints:
+        return F.clamp_keypoints(inpt)  # type: ignore[return-value]

torchvision/transforms/v2/functional/__init__.py

@@ -6,7 +6,7 @@
     clamp_bounding_boxes,
     clamp_keypoints,
     convert_bounding_box_format,
-    convert_bounding_boxes_to_points,
+    convert_bounding_boxes_to_points,  #TODOKP also needs docs
     get_dimensions_image,
     get_dimensions_video,
     get_dimensions,
@@ -157,7 +157,6 @@
     normalize_image,
     normalize_video,
     sanitize_bounding_boxes,
-    sanitize_keypoints,
     to_dtype,
     to_dtype_image,
     to_dtype_video,

torchvision/transforms/v2/functional/_meta.py

@@ -457,12 +457,13 @@ def clamp_bounding_boxes(
 def _clamp_keypoints(keypoints: torch.Tensor, canvas_size: tuple[int, int]) -> torch.Tensor:
     dtype = keypoints.dtype
     keypoints = keypoints.clone() if keypoints.is_floating_point() else keypoints.float()
-    keypoints[..., 0].clamp_(min=0, max=canvas_size[1])
-    keypoints[..., 1].clamp_(min=0, max=canvas_size[0])
+    # Note that max is canvas_size[i] - 1 and not can canvas_size[i] like for
+    # bounding boxes.
+    keypoints[..., 0].clamp_(min=0, max=canvas_size[1] - 1)
+    keypoints[..., 1].clamp_(min=0, max=canvas_size[0] - 1)
     return keypoints.to(dtype=dtype)
 
 
-# TODOKP there is no corresponding transform and this isn't tested
 def clamp_keypoints(
     inpt: torch.Tensor,
     canvas_size: Optional[tuple[int, int]] = None,
@@ -473,7 +474,7 @@ def clamp_keypoints(
     if torch.jit.is_scripting() or is_pure_tensor(inpt):
 
         if canvas_size is None:
-            raise ValueError("For pure tensor inputs, `canvas_size` have to be passed.")
+            raise ValueError("For pure tensor inputs, `canvas_size` has to be passed.")
         return _clamp_keypoints(inpt, canvas_size=canvas_size)
     elif isinstance(inpt, tv_tensors.KeyPoints):
         if canvas_size is not None: