add clamping_mode parameter to clamp_keypoints functional and class

Author: AntoineSimoulin

torchvision/transforms/v2/_meta.py

@@ -2,7 +2,7 @@
 
 from torchvision import tv_tensors
 from torchvision.transforms.v2 import functional as F, Transform
-from torchvision.tv_tensors._bounding_boxes import CLAMPING_MODE_TYPE
+from torchvision.tv_tensors import CLAMPING_MODE_TYPE
 
 
 class ConvertBoundingBoxFormat(Transform):
@@ -46,17 +46,26 @@ def transform(self, inpt: tv_tensors.BoundingBoxes, params: dict[str, Any]) -> t
 class ClampKeyPoints(Transform):
     """Clamp keypoints to their corresponding image dimensions.
 
-    The clamping is done according to the keypoints' ``canvas_size`` meta-data.
+     Args:
+        clamping_mode: Default is "auto" which relies on the input keypoint'
+            ``clamping_mode`` attribute.
+            The clamping is done according to the keypoints' ``canvas_size`` meta-data.
+            Read more in :ref:`clamping_mode_tuto`
+            for more details on how to use this transform.
+            
     """
+    def __init__(self, clamping_mode: Union[CLAMPING_MODE_TYPE, str] = "auto") -> None:
+        super().__init__()
+        self.clamping_mode = clamping_mode
 
     _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]
+        return F.clamp_keypoints(inpt, clamping_mode=self.clamping_mode)  # type: ignore[return-value]
 
 
 class SetClampingMode(Transform):
-    """Sets the ``clamping_mode`` attribute of the bounding boxes for future transforms.
+    """Sets the ``clamping_mode`` attribute of the bounding boxes and keypoints for future transforms.
 
 
 
@@ -73,9 +82,9 @@ def __init__(self, clamping_mode: CLAMPING_MODE_TYPE) -> None:
         if self.clamping_mode not in (None, "soft", "hard"):
             raise ValueError(f"clamping_mode must be soft, hard or None, got {clamping_mode}")
 
-    _transformed_types = (tv_tensors.BoundingBoxes,)
+    _transformed_types = (tv_tensors.BoundingBoxes, tv_tensors.KeyPoints)
 
-    def transform(self, inpt: tv_tensors.BoundingBoxes, params: dict[str, Any]) -> tv_tensors.BoundingBoxes:
-        out: tv_tensors.BoundingBoxes = inpt.clone()  # type: ignore[assignment]
+    def transform(self, inpt: tv_tensors.TVTensor, params: dict[str, Any]) -> tv_tensors.TVTensor:
+        out: tv_tensors.TVTensor = inpt.clone()  # type: ignore[assignment]
         out.clamping_mode = self.clamping_mode
         return out

torchvision/transforms/v2/functional/_geometry.py

@@ -20,7 +20,7 @@
     pil_to_tensor,
     to_pil_image,
 )
-from torchvision.tv_tensors._bounding_boxes import CLAMPING_MODE_TYPE
+from torchvision.tv_tensors import CLAMPING_MODE_TYPE
 
 from torchvision.utils import _log_api_usage_once
 
@@ -67,16 +67,16 @@ def horizontal_flip_mask(mask: torch.Tensor) -> torch.Tensor:
     return horizontal_flip_image(mask)
 
 
-def horizontal_flip_keypoints(keypoints: torch.Tensor, canvas_size: tuple[int, int]):
+def horizontal_flip_keypoints(keypoints: torch.Tensor, canvas_size: tuple[int, int], clamping_mode: CLAMPING_MODE_TYPE = "soft"):
     shape = keypoints.shape
     keypoints = keypoints.clone().reshape(-1, 2)
     keypoints[..., 0] = keypoints[..., 0].sub_(canvas_size[1] - 1).neg_()
-    return clamp_keypoints(keypoints.reshape(shape), canvas_size=canvas_size)
+    return clamp_keypoints(keypoints.reshape(shape), canvas_size=canvas_size, clamping_mode=clamping_mode)
 
 
 @_register_kernel_internal(horizontal_flip, tv_tensors.KeyPoints, tv_tensor_wrapper=False)
 def _horizontal_flip_keypoints_dispatch(keypoints: tv_tensors.KeyPoints):
-    out = horizontal_flip_keypoints(keypoints.as_subclass(torch.Tensor), canvas_size=keypoints.canvas_size)
+    out = horizontal_flip_keypoints(keypoints.as_subclass(torch.Tensor), canvas_size=keypoints.canvas_size, clamping_mode=keypoints.clamping_mode)
     return tv_tensors.wrap(out, like=keypoints)
 
 
@@ -155,11 +155,11 @@ def vertical_flip_mask(mask: torch.Tensor) -> torch.Tensor:
     return vertical_flip_image(mask)
 
 
-def vertical_flip_keypoints(keypoints: torch.Tensor, canvas_size: tuple[int, int]) -> torch.Tensor:
+def vertical_flip_keypoints(keypoints: torch.Tensor, canvas_size: tuple[int, int], clamping_mode: CLAMPING_MODE_TYPE = "soft",) -> torch.Tensor:
     shape = keypoints.shape
     keypoints = keypoints.clone().reshape(-1, 2)
     keypoints[..., 1] = keypoints[..., 1].sub_(canvas_size[0] - 1).neg_()
-    return clamp_keypoints(keypoints.reshape(shape), canvas_size=canvas_size)
+    return clamp_keypoints(keypoints.reshape(shape), canvas_size=canvas_size, clamping_mode=clamping_mode)
 
 
 def vertical_flip_bounding_boxes(
@@ -199,7 +199,7 @@ def vertical_flip_bounding_boxes(
 
 @_register_kernel_internal(vertical_flip, tv_tensors.KeyPoints, tv_tensor_wrapper=False)
 def _vertical_flip_keypoints_dispatch(inpt: tv_tensors.KeyPoints) -> tv_tensors.KeyPoints:
-    output = vertical_flip_keypoints(inpt.as_subclass(torch.Tensor), canvas_size=inpt.canvas_size)
+    output = vertical_flip_keypoints(inpt.as_subclass(torch.Tensor), canvas_size=inpt.canvas_size, clamping_mode=inpt.clamping_mode)
     return tv_tensors.wrap(output, like=inpt)
 
 
@@ -968,6 +968,7 @@ def _affine_keypoints_with_expand(
     shear: list[float],
     center: Optional[list[float]] = None,
     expand: bool = False,
+    clamping_mode: CLAMPING_MODE_TYPE = "soft",
 ) -> tuple[torch.Tensor, tuple[int, int]]:
     if keypoints.numel() == 0:
         return keypoints, canvas_size
@@ -1026,7 +1027,7 @@ def _affine_keypoints_with_expand(
         new_width, new_height = _compute_affine_output_size(affine_vector, width, height)
         canvas_size = (new_height, new_width)
 
-    out_keypoints = clamp_keypoints(transformed_points, canvas_size=canvas_size).reshape(original_shape)
+    out_keypoints = clamp_keypoints(transformed_points, canvas_size=canvas_size, clamping_mode=clamping_mode).reshape(original_shape)
     out_keypoints = out_keypoints.to(original_dtype)
 
     return out_keypoints, canvas_size
@@ -1040,6 +1041,7 @@ def affine_keypoints(
     scale: float,
     shear: list[float],
     center: Optional[list[float]] = None,
+    clamping_mode: CLAMPING_MODE_TYPE = "soft",
 ):
     return _affine_keypoints_with_expand(
         keypoints=keypoints,
@@ -1050,6 +1052,7 @@ def affine_keypoints(
         shear=shear,
         center=center,
         expand=False,
+        clamping_mode=clamping_mode,
     )
 
 
@@ -1071,6 +1074,7 @@ def _affine_keypoints_dispatch(
         scale=scale,
         shear=shear,
         center=center,
+        clamping_mode=inpt.clamping_mode,
     )
     return tv_tensors.wrap(output, like=inpt, canvas_size=canvas_size)
 
@@ -1393,6 +1397,7 @@ def rotate_keypoints(
     angle: float,
     expand: bool = False,
     center: Optional[list[float]] = None,
+    clamping_mode: CLAMPING_MODE_TYPE = "soft",
 ) -> tuple[torch.Tensor, tuple[int, int]]:
     return _affine_keypoints_with_expand(
         keypoints=keypoints,
@@ -1403,6 +1408,7 @@ def rotate_keypoints(
         shear=[0.0, 0.0],
         center=center,
         expand=expand,
+        clamping_mode=clamping_mode,
     )
 
 
@@ -1411,7 +1417,7 @@ def _rotate_keypoints_dispatch(
     inpt: tv_tensors.KeyPoints, angle: float, expand: bool = False, center: Optional[list[float]] = None, **kwargs
 ) -> tv_tensors.KeyPoints:
     output, canvas_size = rotate_keypoints(
-        inpt, canvas_size=inpt.canvas_size, angle=angle, center=center, expand=expand
+        inpt, canvas_size=inpt.canvas_size, angle=angle, center=center, expand=expand, clamping_mode=inpt.clamping_mode,
     )
     return tv_tensors.wrap(output, like=inpt, canvas_size=canvas_size)
 
@@ -1683,7 +1689,7 @@ def pad_mask(
 
 
 def pad_keypoints(
-    keypoints: torch.Tensor, canvas_size: tuple[int, int], padding: list[int], padding_mode: str = "constant"
+    keypoints: torch.Tensor, canvas_size: tuple[int, int], padding: list[int], padding_mode: str = "constant", clamping_mode: CLAMPING_MODE_TYPE = "soft"
 ):
     SUPPORTED_MODES = ["constant"]
     if padding_mode not in SUPPORTED_MODES:
@@ -1695,20 +1701,21 @@ def pad_keypoints(
     left, right, top, bottom = _parse_pad_padding(padding)
     pad = torch.tensor([left, top], dtype=keypoints.dtype, device=keypoints.device)
     canvas_size = (canvas_size[0] + top + bottom, canvas_size[1] + left + right)
-    return clamp_keypoints(keypoints + pad, canvas_size), canvas_size
+    return clamp_keypoints(keypoints + pad, canvas_size=canvas_size, clamping_mode=clamping_mode), canvas_size
 
 
 @_register_kernel_internal(pad, tv_tensors.KeyPoints, tv_tensor_wrapper=False)
 def _pad_keypoints_dispatch(
-    keypoints: tv_tensors.KeyPoints, padding: list[int], padding_mode: str = "constant", **kwargs
+    inpt: tv_tensors.KeyPoints, padding: list[int], padding_mode: str = "constant", **kwargs
 ) -> tv_tensors.KeyPoints:
     output, canvas_size = pad_keypoints(
-        keypoints.as_subclass(torch.Tensor),
-        canvas_size=keypoints.canvas_size,
+        inpt.as_subclass(torch.Tensor),
+        canvas_size=inpt.canvas_size,
         padding=padding,
         padding_mode=padding_mode,
+        clamping_mode=inpt.clamping_mode,
     )
-    return tv_tensors.wrap(output, like=keypoints, canvas_size=canvas_size)
+    return tv_tensors.wrap(output, like=inpt, canvas_size=canvas_size)
 
 
 def pad_bounding_boxes(
@@ -1812,19 +1819,20 @@ def crop_keypoints(
     left: int,
     height: int,
     width: int,
+    clamping_mode: CLAMPING_MODE_TYPE = "soft",
 ) -> tuple[torch.Tensor, tuple[int, int]]:
 
     keypoints = keypoints - torch.tensor([left, top], dtype=keypoints.dtype, device=keypoints.device)
     canvas_size = (height, width)
 
-    return clamp_keypoints(keypoints, canvas_size=canvas_size), canvas_size
+    return clamp_keypoints(keypoints, canvas_size=canvas_size, clamping_mode=clamping_mode), canvas_size
 
 
 @_register_kernel_internal(crop, tv_tensors.KeyPoints, tv_tensor_wrapper=False)
 def _crop_keypoints_dispatch(
     inpt: tv_tensors.KeyPoints, top: int, left: int, height: int, width: int
 ) -> tv_tensors.KeyPoints:
-    output, canvas_size = crop_keypoints(inpt.as_subclass(torch.Tensor), top=top, left=left, height=height, width=width)
+    output, canvas_size = crop_keypoints(inpt.as_subclass(torch.Tensor), top=top, left=left, height=height, width=width, clamping_mode=inpt.clamping_mode)
     return tv_tensors.wrap(output, like=inpt, canvas_size=canvas_size)
 
 
@@ -2024,6 +2032,7 @@ def perspective_keypoints(
     startpoints: Optional[list[list[int]]],
     endpoints: Optional[list[list[int]]],
     coefficients: Optional[list[float]] = None,
+    clamping_mode: CLAMPING_MODE_TYPE = "soft",
 ):
     if keypoints.numel() == 0:
         return keypoints
@@ -2047,7 +2056,7 @@ def perspective_keypoints(
     numer_points = torch.matmul(points, theta1.T)
     denom_points = torch.matmul(points, theta2.T)
     transformed_points = numer_points.div_(denom_points)
-    return clamp_keypoints(transformed_points.to(keypoints.dtype), canvas_size).reshape(original_shape)
+    return clamp_keypoints(transformed_points.to(keypoints.dtype), canvas_size=canvas_size, clamping_mode=clamping_mode).reshape(original_shape)
 
 
 @_register_kernel_internal(perspective, tv_tensors.KeyPoints, tv_tensor_wrapper=False)
@@ -2064,6 +2073,7 @@ def _perspective_keypoints_dispatch(
         startpoints=startpoints,
         endpoints=endpoints,
         coefficients=coefficients,
+        clamping_mode=inpt.clamping_mode,
     )
     return tv_tensors.wrap(output, like=inpt)
 
@@ -2344,7 +2354,7 @@ def _create_identity_grid(size: tuple[int, int], device: torch.device, dtype: to
 
 
 def elastic_keypoints(
-    keypoints: torch.Tensor, canvas_size: tuple[int, int], displacement: torch.Tensor
+    keypoints: torch.Tensor, canvas_size: tuple[int, int], displacement: torch.Tensor, clamping_mode: CLAMPING_MODE_TYPE = "soft",
 ) -> torch.Tensor:
     expected_shape = (1, canvas_size[0], canvas_size[1], 2)
     if not isinstance(displacement, torch.Tensor):
@@ -2376,12 +2386,12 @@ def elastic_keypoints(
     t_size = torch.tensor(canvas_size[::-1], device=displacement.device, dtype=displacement.dtype)
     transformed_points = inv_grid[0, index_y, index_x, :].add_(1).mul_(0.5 * t_size).sub_(0.5)
 
-    return clamp_keypoints(transformed_points.to(keypoints.dtype), canvas_size=canvas_size).reshape(original_shape)
+    return clamp_keypoints(transformed_points.to(keypoints.dtype), canvas_size=canvas_size, clamping_mode=clamping_mode).reshape(original_shape)
 
 
 @_register_kernel_internal(elastic, tv_tensors.KeyPoints, tv_tensor_wrapper=False)
 def _elastic_keypoints_dispatch(inpt: tv_tensors.KeyPoints, displacement: torch.Tensor, **kwargs):
-    output = elastic_keypoints(inpt.as_subclass(torch.Tensor), canvas_size=inpt.canvas_size, displacement=displacement)
+    output = elastic_keypoints(inpt.as_subclass(torch.Tensor), canvas_size=inpt.canvas_size, displacement=displacement, clamping_mode=inpt.clamping_mode)
     return tv_tensors.wrap(output, like=inpt)
 
 
@@ -2578,16 +2588,16 @@ def _center_crop_image_pil(image: PIL.Image.Image, output_size: list[int]) -> PI
     return _crop_image_pil(image, crop_top, crop_left, crop_height, crop_width)
 
 
-def center_crop_keypoints(inpt: torch.Tensor, canvas_size: tuple[int, int], output_size: list[int]):
+def center_crop_keypoints(inpt: torch.Tensor, canvas_size: tuple[int, int], output_size: list[int], clamping_mode: CLAMPING_MODE_TYPE = "soft",):
     crop_height, crop_width = _center_crop_parse_output_size(output_size)
     crop_top, crop_left = _center_crop_compute_crop_anchor(crop_height, crop_width, *canvas_size)
-    return crop_keypoints(inpt, top=crop_top, left=crop_left, height=crop_height, width=crop_width)
+    return crop_keypoints(inpt, top=crop_top, left=crop_left, height=crop_height, width=crop_width, clamping_mode=clamping_mode)
 
 
 @_register_kernel_internal(center_crop, tv_tensors.KeyPoints, tv_tensor_wrapper=False)
 def _center_crop_keypoints_dispatch(inpt: tv_tensors.KeyPoints, output_size: list[int]) -> tv_tensors.KeyPoints:
     output, canvas_size = center_crop_keypoints(
-        inpt.as_subclass(torch.Tensor), canvas_size=inpt.canvas_size, output_size=output_size
+        inpt.as_subclass(torch.Tensor), canvas_size=inpt.canvas_size, output_size=output_size, clamping_mode=inpt.clamping_mode,
     )
     return tv_tensors.wrap(output, like=inpt, canvas_size=canvas_size)
 
@@ -2745,8 +2755,9 @@ def resized_crop_keypoints(
     height: int,
     width: int,
     size: list[int],
+    clamping_mode: CLAMPING_MODE_TYPE = "soft",
 ) -> tuple[torch.Tensor, tuple[int, int]]:
-    keypoints, canvas_size = crop_keypoints(keypoints, top, left, height, width)
+    keypoints, canvas_size = crop_keypoints(keypoints, top, left, height, width, clamping_mode=clamping_mode)
     return resize_keypoints(keypoints, size=size, canvas_size=canvas_size)
 
 
@@ -2755,7 +2766,7 @@ def _resized_crop_keypoints_dispatch(
     inpt: tv_tensors.BoundingBoxes, top: int, left: int, height: int, width: int, size: list[int], **kwargs
 ):
     output, canvas_size = resized_crop_keypoints(
-        inpt.as_subclass(torch.Tensor), top=top, left=left, height=height, width=width, size=size
+        inpt.as_subclass(torch.Tensor), top=top, left=left, height=height, width=width, size=size, clamping_mode=inpt.clamping_mode,
     )
     return tv_tensors.wrap(output, like=inpt, canvas_size=canvas_size)
 

torchvision/transforms/v2/functional/_meta.py

@@ -5,7 +5,7 @@
 from torchvision import tv_tensors
 from torchvision.transforms import _functional_pil as _FP
 from torchvision.tv_tensors import BoundingBoxFormat
-from torchvision.tv_tensors._bounding_boxes import CLAMPING_MODE_TYPE
+from torchvision.tv_tensors import CLAMPING_MODE_TYPE
 
 from torchvision.utils import _log_api_usage_once
 
@@ -653,7 +653,9 @@ def clamp_bounding_boxes(
         )
 
 
-def _clamp_keypoints(keypoints: torch.Tensor, canvas_size: tuple[int, int]) -> torch.Tensor:
+def _clamp_keypoints(keypoints: torch.Tensor, canvas_size: tuple[int, int], clamping_mode: CLAMPING_MODE_TYPE) -> torch.Tensor:
+    if clamping_mode is None or clamping_mode != "hard":
+        return keypoints.clone()
     dtype = keypoints.dtype
     keypoints = keypoints.clone() if keypoints.is_floating_point() else keypoints.float()
     # Note that max is canvas_size[i] - 1 and not can canvas_size[i] like for
@@ -666,20 +668,26 @@ def _clamp_keypoints(keypoints: torch.Tensor, canvas_size: tuple[int, int]) -> t
 def clamp_keypoints(
     inpt: torch.Tensor,
     canvas_size: Optional[tuple[int, int]] = None,
+    clamping_mode: Union[CLAMPING_MODE_TYPE, str] = "auto",
 ) -> torch.Tensor:
     """See :func:`~torchvision.transforms.v2.ClampKeyPoints` for details."""
     if not torch.jit.is_scripting():
         _log_api_usage_once(clamp_keypoints)
 
+    if clamping_mode is not None and clamping_mode not in ("soft", "hard", "auto"):
+        raise ValueError(f"clamping_mode must be soft, hard, auto or None, got {clamping_mode}")
+
     if torch.jit.is_scripting() or is_pure_tensor(inpt):
 
-        if canvas_size is None:
-            raise ValueError("For pure tensor inputs, `canvas_size` has to be passed.")
-        return _clamp_keypoints(inpt, canvas_size=canvas_size)
+        if canvas_size is None or (clamping_mode is not None and clamping_mode == "auto"):
+            raise ValueError("For pure tensor inputs, `canvas_size` and `clamping_mode` have to be passed.")
+        return _clamp_keypoints(inpt, canvas_size=canvas_size, clamping_mode=clamping_mode)
     elif isinstance(inpt, tv_tensors.KeyPoints):
         if canvas_size is not None:
             raise ValueError("For keypoints tv_tensor inputs, `canvas_size` must not be passed.")
-        output = _clamp_keypoints(inpt.as_subclass(torch.Tensor), canvas_size=inpt.canvas_size)
+        if clamping_mode is None and clamping_mode == "auto":
+            clamping_mode = inpt.clamping_mode
+        output = _clamp_keypoints(inpt.as_subclass(torch.Tensor), canvas_size=inpt.canvas_size, clamping_mode=clamping_mode)
         return tv_tensors.wrap(output, like=inpt)
     else:
         raise TypeError(f"Input can either be a plain tensor or a keypoints tv_tensor, but got {type(inpt)} instead.")