lint

Author: AntoineSimoulin

test/test_transforms_v2.py

@@ -661,7 +661,8 @@ def affine_keypoints(keypoints):
 
     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, clamping_mode=clamping_mode
+        canvas_size=canvas_size,
+        clamping_mode=clamping_mode,
     )
 
 
@@ -5362,11 +5363,9 @@ def perspective_keypoints(keypoints):
             )
 
             # It is important to clamp before casting, especially for CXCYWH format, dtype=int64
-            return F.clamp_keypoints(
-                output,
-                canvas_size=canvas_size,
-                clamping_mode=clamping_mode
-            ).to(dtype=dtype, device=device)
+            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(
@@ -5791,9 +5790,14 @@ def test_clamping_mode(self, constructor_clamping_mode, clamping_mode, pass_pure
             return
 
         keypoints = tv_tensors.KeyPoints(
-            [[0, 100], [0, 100]],canvas_size=(10, 10), clamping_mode=constructor_clamping_mode
+            [[0, 100], [0, 100]], canvas_size=(10, 10), clamping_mode=constructor_clamping_mode
+        )
+        expected_clamped_output = (
+            torch.tensor([[0, 10], [0, 10]]) if clamping_mode == "hard" else torch.tensor([[0, 100], [0, 100]])
+        )
+        expected_clamped_output = (
+            torch.tensor([[0, 9], [0, 9]]) if clamping_mode == "hard" else torch.tensor([[0, 100], [0, 100]])
         )
-        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(

torchvision/transforms/v2/_meta.py

@@ -46,14 +46,15 @@ def transform(self, inpt: tv_tensors.BoundingBoxes, params: dict[str, Any]) -> t
 class ClampKeyPoints(Transform):
     """Clamp keypoints to their corresponding image dimensions.
 
-     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.
-            
+    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

torchvision/transforms/v2/functional/_geometry.py

@@ -67,7 +67,9 @@ 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], clamping_mode: CLAMPING_MODE_TYPE = "soft"):
+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_()
@@ -76,7 +78,9 @@ def horizontal_flip_keypoints(keypoints: torch.Tensor, canvas_size: tuple[int, i
 
 @_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, clamping_mode=keypoints.clamping_mode)
+    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,7 +159,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], clamping_mode: CLAMPING_MODE_TYPE = "soft",) -> 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_()
@@ -199,7 +207,9 @@ 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, clamping_mode=inpt.clamping_mode)
+    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)
 
 
@@ -1027,7 +1037,9 @@ 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, clamping_mode=clamping_mode).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
@@ -1417,7 +1429,12 @@ 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, clamping_mode=inpt.clamping_mode,
+        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)
 
@@ -1689,7 +1706,11 @@ def pad_mask(
 
 
 def pad_keypoints(
-    keypoints: torch.Tensor, canvas_size: tuple[int, int], padding: list[int], padding_mode: str = "constant", clamping_mode: CLAMPING_MODE_TYPE = "soft"
+    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:
@@ -1832,7 +1853,9 @@ def crop_keypoints(
 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, clamping_mode=inpt.clamping_mode)
+    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)
 
 
@@ -2056,7 +2079,9 @@ 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=canvas_size, clamping_mode=clamping_mode).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)
@@ -2354,7 +2379,10 @@ 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, clamping_mode: CLAMPING_MODE_TYPE = "soft",
+    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):
@@ -2386,12 +2414,19 @@ 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, clamping_mode=clamping_mode).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, clamping_mode=inpt.clamping_mode)
+    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)
 
 
@@ -2588,16 +2623,26 @@ 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], clamping_mode: CLAMPING_MODE_TYPE = "soft",):
+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, clamping_mode=clamping_mode)
+    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, clamping_mode=inpt.clamping_mode,
+        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)
 
@@ -2766,7 +2811,13 @@ 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, clamping_mode=inpt.clamping_mode,
+        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

@@ -4,8 +4,7 @@
 import torch
 from torchvision import tv_tensors
 from torchvision.transforms import _functional_pil as _FP
-from torchvision.tv_tensors import BoundingBoxFormat
-from torchvision.tv_tensors import CLAMPING_MODE_TYPE
+from torchvision.tv_tensors import BoundingBoxFormat, CLAMPING_MODE_TYPE
 
 from torchvision.utils import _log_api_usage_once
 
@@ -653,7 +652,9 @@ def clamp_bounding_boxes(
         )
 
 
-def _clamp_keypoints(keypoints: torch.Tensor, canvas_size: tuple[int, int], clamping_mode: CLAMPING_MODE_TYPE) -> 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
@@ -687,7 +688,9 @@ def clamp_keypoints(
             raise ValueError("For keypoints tv_tensor inputs, `canvas_size` must not be passed.")
         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)
+        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.")

torchvision/tv_tensors/__init__.py

@@ -1,6 +1,6 @@
 import torch
 
-from ._bounding_boxes import BoundingBoxes, BoundingBoxFormat, is_rotated_bounding_format, CLAMPING_MODE_TYPE
+from ._bounding_boxes import BoundingBoxes, BoundingBoxFormat, CLAMPING_MODE_TYPE, is_rotated_bounding_format
 from ._image import Image
 from ._keypoints import KeyPoints
 from ._mask import Mask

torchvision/tv_tensors/_keypoints.py

@@ -5,9 +5,10 @@
 import torch
 from torch.utils._pytree import tree_flatten
 
-from ._tv_tensor import TVTensor
 from ._bounding_boxes import CLAMPING_MODE_TYPE
 
+from ._tv_tensor import TVTensor
+
 
 class KeyPoints(TVTensor):
     """:class:`torch.Tensor` subclass for tensors with shape ``[..., 2]`` that represent points in an image.
@@ -103,7 +104,10 @@ def _wrap_output(
             output = KeyPoints._wrap(output, canvas_size=canvas_size, clamping_mode=clamping_mode, check_dims=False)
         elif isinstance(output, (tuple, list)):
             # This branch exists for chunk() and unbind()
-            output = type(output)(KeyPoints._wrap(part, canvas_size=canvas_size, clamping_mode=clamping_mode, check_dims=False) for part in output)
+            output = type(output)(
+                KeyPoints._wrap(part, canvas_size=canvas_size, clamping_mode=clamping_mode, check_dims=False)
+                for part in output
+            )
         return output
 
     def __repr__(self, *, tensor_contents: Any = None) -> str:  # type: ignore[override]