More tests

Author: NicolasHug

test/test_transforms_v2.py

@@ -1461,7 +1461,7 @@ def test_kernel_video(self):
 
     @pytest.mark.parametrize(
         "make_input",
-        [make_image_tensor, make_image_pil, make_image, make_bounding_boxes, make_segmentation_mask, make_video],
+        [make_image_tensor, make_image_pil, make_image, make_bounding_boxes, make_segmentation_mask, make_video, make_keypoints],
     )
     def test_functional(self, make_input):
         check_functional(F.affine, make_input(), **self._MINIMAL_AFFINE_KWARGS)
@@ -1483,7 +1483,7 @@ def test_functional_signature(self, kernel, input_type):
 
     @pytest.mark.parametrize(
         "make_input",
-        [make_image_tensor, make_image_pil, make_image, make_bounding_boxes, make_segmentation_mask, make_video],
+        [make_image_tensor, make_image_pil, make_image, make_bounding_boxes, make_segmentation_mask, make_video, make_keypoints],
     )
     @pytest.mark.parametrize("device", cpu_and_cuda())
     def test_transform(self, make_input, device):
@@ -1588,17 +1588,6 @@ def _reference_affine_bounding_boxes(self, bounding_boxes, *, angle, translate,
             ),
         )
 
-    def _reference_affine_keypoints(self, keypoints, *, angle, translate, scale, shear, center):
-        if center is None:
-            center = [s * 0.5 for s in keypoints.canvas_size[::-1]]
-
-        return reference_affine_keypoints_helper(
-            keypoints,
-            affine_matrix=self._compute_affine_matrix(
-                angle=angle, translate=translate, scale=scale, shear=shear, center=center
-            ),
-        )
-
     @pytest.mark.parametrize("format", SUPPORTED_BOX_FORMATS)
     @pytest.mark.parametrize("angle", _CORRECTNESS_AFFINE_KWARGS["angle"])
     @pytest.mark.parametrize("translate", _CORRECTNESS_AFFINE_KWARGS["translate"])
@@ -1645,6 +1634,18 @@ def test_transform_bounding_boxes_correctness(self, format, center, seed):
 
         torch.testing.assert_close(actual, expected)
 
+    def _reference_affine_keypoints(self, keypoints, *, angle, translate, scale, shear, center):
+        if center is None:
+            center = [s * 0.5 for s in keypoints.canvas_size[::-1]]
+
+        return reference_affine_keypoints_helper(
+            keypoints,
+            affine_matrix=self._compute_affine_matrix(
+                angle=angle, translate=translate, scale=scale, shear=shear, center=center
+            ),
+        )
+
+
     @pytest.mark.parametrize("angle", _CORRECTNESS_AFFINE_KWARGS["angle"])
     @pytest.mark.parametrize("translate", _CORRECTNESS_AFFINE_KWARGS["translate"])
     @pytest.mark.parametrize("scale", _CORRECTNESS_AFFINE_KWARGS["scale"])

torchvision/transforms/v2/functional/_geometry.py

@@ -868,6 +868,7 @@ def _affine_image_pil(
     return _FP.affine(image, matrix, interpolation=pil_modes_mapping[interpolation], fill=fill)
 
 
+# TODO: Consider merging/unifying this with the bbox implementation
 def _affine_keypoints_with_expand(
     keypoints: torch.Tensor,
     canvas_size: tuple[int, int],
@@ -935,10 +936,10 @@ def _affine_keypoints_with_expand(
         new_width, new_height = _compute_affine_output_size(affine_vector, width, height)
         canvas_size = (new_height, new_width)
 
-    out_kkpoints = clamp_keypoints(transformed_points, canvas_size=canvas_size).reshape(original_shape)
-    out_kkpoints = out_kkpoints.to(original_dtype)
+    out_keypoints = clamp_keypoints(transformed_points, canvas_size=canvas_size).reshape(original_shape)
+    out_keypoints = out_keypoints.to(original_dtype)
 
-    return out_kkpoints, canvas_size
+    return out_keypoints, canvas_size
 
 
 def affine_keypoints(