Simplify _get_sanitize_keypoints_mask

Author: AntoineSimoulin

test/test_transforms_v2.py

@@ -7401,24 +7401,22 @@ class TestSanitizeKeyPoints:
     def _make_keypoints_with_validity(
         self,
         canvas_size=(100, 100),
-        min_valid_edge_distance=0,
-        min_invalid_points=1,
         shape="2d",  # "2d", "3d", "4d" for different keypoint shapes
     ):
         """Create keypoints with known validity for testing."""
         canvas_h, canvas_w = canvas_size
 
         if shape == "2d":  # [N_points, 2]
             keypoints_data = [
-                ([5, 5], min_valid_edge_distance <= 5),  # Valid point inside image
-                ([canvas_w - 6, canvas_h - 6], min_valid_edge_distance <= 5),  # Valid point near corner
+                ([5, 5], True),  # Valid point inside image
+                ([canvas_w - 6, canvas_h - 6], True),  # Valid point near corner
                 ([canvas_w // 2, canvas_h // 2], True),  # Valid point in center
                 ([-1, canvas_h // 2], False),  # Invalid: x < 0
                 ([canvas_w // 2, -1], False),  # Invalid: y < 0
                 ([canvas_w, canvas_h // 2], False),  # Invalid: x >= canvas_w
                 ([canvas_w // 2, canvas_h], False),  # Invalid: y >= canvas_h
-                ([0, 0], min_valid_edge_distance <= 0),  # Edge case: exactly on edge
-                ([canvas_w - 1, canvas_h - 1], min_valid_edge_distance <= 0),  # Edge case: exactly on edge
+                ([0, 0], True),  # Edge case: exactly on edge
+                ([canvas_w - 1, canvas_h - 1], True),  # Edge case: exactly on edge
             ]
             points, validity = zip(*keypoints_data)
             keypoints = torch.tensor(points, dtype=torch.float32)
@@ -7429,11 +7427,11 @@ def _make_keypoints_with_validity(
                 # Group 1: All points valid
                 ([[10, 10], [20, 20], [30, 30]], True),
                 # Group 2: One invalid point (should be removed if min_invalid_points=1)
-                ([[10, 10], [20, 20], [-5, 30]], min_invalid_points > 1),
+                ([[10, 10], [20, 20], [-5, 30]], False),
                 # Group 3: All points invalid
                 ([[-1, -1], [-2, -2], [-3, -3]], False),
                 # Group 4: Mix of valid and invalid (depends on min_invalid_points)
-                ([[10, 10], [-1, 20], [-2, 30]], min_invalid_points > 2),
+                ([[10, 10], [-1, 20], [-2, 30]], False),
             ]
             groups, validity = zip(*keypoints_data)
             keypoints = torch.tensor(groups, dtype=torch.float32)
@@ -7444,7 +7442,7 @@ def _make_keypoints_with_validity(
                 # Object 1: All bones valid
                 ([[[10, 10], [15, 15]], [[20, 20], [25, 25]]], True),
                 # Object 2: One bone with invalid point
-                ([[[10, 10], [15, 15]], [[-1, 20], [25, 25]]], min_invalid_points > 1),
+                ([[[10, 10], [15, 15]], [[-1, 20], [25, 25]]], False),
                 # Object 3: All bones invalid
                 ([[[-1, -1], [-2, -2]], [[-3, -3], [-4, -4]]], False),
             ]
@@ -7457,26 +7455,14 @@ def _make_keypoints_with_validity(
         return keypoints, validity
 
     @pytest.mark.parametrize("shape", ["2d", "3d", "4d"])
-    @pytest.mark.parametrize("min_valid_edge_distance", [0, 1, 5, 6])
-    @pytest.mark.parametrize("min_invalid_points", [1, 2, 0.5])
     @pytest.mark.parametrize("input_type", [torch.Tensor, tv_tensors.KeyPoints])
-    def test_functional(self, shape, min_valid_edge_distance, min_invalid_points, input_type):
+    def test_functional(self, shape, input_type):
         """Test the sanitize_keypoints functional interface."""
-        # Check for invalid configuration
-        if shape == "2d" and min_invalid_points > 1:
-            pytest.xfail("min_invalid_points > 1 does not make sense for 2D keypoints")
 
         # Create inputs
         canvas_size = (50, 50)
-        if isinstance(min_invalid_points, float):
-            num_groups = 4 if shape == "4d" else 3
-            min_invalid_points_int = math.ceil(min_invalid_points * num_groups)
-        else:
-            min_invalid_points_int = min_invalid_points
         keypoints, expected_validity = self._make_keypoints_with_validity(
             canvas_size=canvas_size,
-            min_valid_edge_distance=min_valid_edge_distance,
-            min_invalid_points=min_invalid_points_int,
             shape=shape,
         )
 
@@ -7490,8 +7476,6 @@ def test_functional(self, shape, min_valid_edge_distance, min_invalid_points, in
         result_keypoints, valid_mask = F.sanitize_keypoints(
             keypoints,
             canvas_size=canvas_size_arg,
-            min_valid_edge_distance=min_valid_edge_distance,
-            min_invalid_points=min_invalid_points,
         )
 
         # Check return types
@@ -7522,8 +7506,6 @@ def test_kernel(self, shape):
         )
 
     @pytest.mark.parametrize("shape", ["2d", "3d", "4d"])
-    @pytest.mark.parametrize("min_valid_edge_distance", [0, 2])
-    @pytest.mark.parametrize("min_invalid_points", [1, 0.3])
     @pytest.mark.parametrize(
         "labels_getter",
         (
@@ -7536,26 +7518,15 @@ def test_kernel(self, shape):
         ),
     )
     @pytest.mark.parametrize("sample_type", (tuple, dict))
-    def test_transform(self, shape, min_valid_edge_distance, min_invalid_points, labels_getter, sample_type):
+    def test_transform(self, shape, labels_getter, sample_type):
         """Test the SanitizeKeyPoints transform class."""
         if sample_type is tuple and not isinstance(labels_getter, str):
             # Lambda-based labels_getter doesn't work with tuple input
             return
 
-        # Check for invalid configuration
-        if shape == "2d" and min_invalid_points > 1:
-            pytest.xfail("min_invalid_points > 1 does not make sense for 2D keypoints")
-
         canvas_size = (40, 40)
-        if isinstance(min_invalid_points, float):
-            num_groups = 4 if shape == "4d" else 3
-            min_invalid_points_int = math.ceil(min_invalid_points * num_groups)
-        else:
-            min_invalid_points_int = min_invalid_points
         keypoints, expected_validity = self._make_keypoints_with_validity(
             canvas_size=canvas_size,
-            min_valid_edge_distance=min_valid_edge_distance,
-            min_invalid_points=min_invalid_points_int,
             shape=shape,
         )
 
@@ -7585,8 +7556,6 @@ def test_transform(self, shape, min_valid_edge_distance, min_invalid_points, lab
 
         # Apply transform
         transform = transforms.SanitizeKeyPoints(
-            min_valid_edge_distance=min_valid_edge_distance,
-            min_invalid_points=min_invalid_points,
             labels_getter=labels_getter,
         )
         out = transform(sample)
@@ -7644,6 +7613,7 @@ def test_edge_cases(self):
         # Test empty keypoints
         empty_keypoints = tv_tensors.KeyPoints(torch.empty(0, 2), canvas_size=canvas_size)
         result, valid_mask = F.sanitize_keypoints(empty_keypoints)
+        print(empty_keypoints, result, valid_mask)
         assert tuple(result.shape) == (0, 2)
         assert valid_mask.shape[0] == 0
 
@@ -7659,43 +7629,6 @@ def test_edge_cases(self):
         assert tuple(result.shape) == (0, 2)
         assert not valid_mask.any()
 
-    def test_min_invalid_points_fraction(self):
-        """Test min_invalid_points as a fraction."""
-        canvas_size = (20, 20)
-
-        # Create 3D keypoints with 4 points per object
-        keypoints = torch.tensor(
-            [
-                # Object 1: 1 invalid point out of 4 (25% invalid)
-                [[5, 5], [10, 10], [15, 15], [-1, -1]],
-                # Object 2: 2 invalid points out of 4 (50% invalid)
-                [[5, 5], [10, 10], [-1, -1], [-2, -2]],
-                # Object 3: 3 invalid points out of 4 (75% invalid)
-                [[5, 5], [-1, -1], [-2, -2], [-3, -3]],
-            ],
-            dtype=torch.float32,
-        )
-
-        keypoints = tv_tensors.KeyPoints(keypoints, canvas_size=canvas_size)
-
-        # Test with 30% threshold - should keep object 1
-        result, valid_mask = F.sanitize_keypoints(keypoints, min_invalid_points=0.3)
-        expected_valid = torch.tensor([True, False, False])
-        assert_equal(valid_mask, expected_valid)
-        assert result.shape[0] == 1
-
-        # Test with 60% threshold - should keep objects 1 and 2
-        result, valid_mask = F.sanitize_keypoints(keypoints, min_invalid_points=0.6)
-        expected_valid = torch.tensor([True, True, False])
-        assert_equal(valid_mask, expected_valid)
-        assert result.shape[0] == 2
-
-        # Test with 100% threshold - should keep all objects
-        result, valid_mask = F.sanitize_keypoints(keypoints, min_invalid_points=1.0)
-        expected_valid = torch.tensor([True, True, True])
-        assert_equal(valid_mask, expected_valid)
-        assert result.shape[0] == 3
-
     def test_errors_functional(self):
         """Test error conditions for the functional interface."""
         good_keypoints = tv_tensors.KeyPoints([[5, 5]], canvas_size=(10, 10))
@@ -7708,16 +7641,6 @@ def test_errors_functional(self):
         with pytest.raises(ValueError, match="canvas_size must be None"):
             F.sanitize_keypoints(good_keypoints, canvas_size=(10, 10))
 
-        # Test invalid min_invalid_points
-        with pytest.raises(ValueError, match="min_invalid_points must be > 0"):
-            F.sanitize_keypoints(good_keypoints, min_invalid_points=0)
-
-        with pytest.raises(ValueError, match="min_invalid_points must be > 0"):
-            F.sanitize_keypoints(good_keypoints, min_invalid_points=-1)
-
-        with pytest.raises(ValueError, match="so min_invalid_points must be 1"):
-            F.sanitize_keypoints(good_keypoints, min_invalid_points=2)
-
     def test_errors_transform(self):
         """Test error conditions for the transform class."""
         good_keypoints = tv_tensors.KeyPoints([[5, 5]], canvas_size=(10, 10))
@@ -7726,10 +7649,6 @@ def test_errors_transform(self):
         with pytest.raises(ValueError, match="labels_getter should either be"):
             transforms.SanitizeKeyPoints(labels_getter="invalid_type")  # type: ignore
 
-        # Test invalid min_invalid_points
-        with pytest.raises(ValueError, match="min_invalid_points must be > 0"):
-            transforms.SanitizeKeyPoints(min_invalid_points=0)
-
         # Test missing labels key
         with pytest.raises(ValueError, match="Could not infer where the labels are"):
             bad_sample = {"keypoints": good_keypoints, "BAD_KEY": torch.tensor([0])}
@@ -7745,10 +7664,6 @@ def test_errors_transform(self):
             bad_sample = {"keypoints": good_keypoints, "labels": torch.tensor([0, 1, 2])}
             transforms.SanitizeKeyPoints(labels_getter="default")(bad_sample)
 
-        # Test min_invalid_points > 1 for 2D keypoints
-        with pytest.raises(ValueError, match="so min_invalid_points must be 1"):
-            transforms.SanitizeKeyPoints(min_invalid_points=2)(good_keypoints)
-
     def test_no_label(self):
         """Test transform without labels."""
         img = make_image()

torchvision/transforms/v2/_misc.py

@@ -473,34 +473,19 @@ class SanitizeKeyPoints(Transform):
     """Remove keypoints outside of the image area and their corresponding labels (if any).
 
     This transform removes keypoints or groups of keypoints and their associated labels that
-    have coordinates outside of their corresponding image or within ``min_valid_edge_distance`` pixels
-    from the image edges.
+    have coordinates outside of their corresponding image.
     If you would instead like to clamp such keypoints to the image edges, use
     :class:`~torchvision.transforms.v2.ClampKeyPoints`.
 
     It is recommended to call it at the end of a pipeline, before passing the
     input to the models.
 
-    Keypoints can be passed as a set of individual keypoints of shape ``[N_points, 2]`` or as a
-    set of objects (e.g., polygons or polygonal chains) consisting of a fixed number of keypoints
-    of shape ``[N_objects, ..., 2]``.
+    Keypoints can be passed as a set of individual keypoints or as a set of objects
+    (e.g., polygons or polygonal chains) consisting of a fixed number of keypoints of shape ``[..., 2]``.
     When groups of keypoints are passed (i.e., an at least 3-dimensional tensor), this transform
     will only remove entire groups, not individual keypoints within a group.
 
     Args:
-        min_valid_edge_distance (int, optional): The minimum distance that keypoints need to be away from the closest image
-            edge along any axis in order to be considered valid. For example, setting this to 0 will only
-            invalidate/remove keypoints outside of the image area, while a value of 1 will also remove keypoints
-            lying exactly on the edge.
-            Default is 0.
-        min_invalid_points (int or float, optional): Minimum number or fraction of invalid keypoints required
-            for a group of keypoints to be removed. For example, setting this to 1 will remove a group of keypoints
-            if any of its keypoints is invalid, while setting it to 2 will only remove groups with at least 2 invalid keypoints.
-            If a float in ``(0.0, 1.0]`` is passed, it represents a fraction of the total number of keypoints in
-            the group. For example, setting this to 0.3 will remove groups of keypoints with at least 30% invalid keypoints.
-            Note that a value of `1` (integer) is very different from `1.0` (float). The former will remove groups
-            with any invalid keypoint, while the latter will only remove groups where all keypoints are invalid.
-            Default is 1.
         labels_getter (callable or str or None, optional): indicates how to identify the labels in the input
             (or anything else that needs to be sanitized along with the keypoints).
             If set to the string ``"default"``, this will try to find a "labels" key in the input (case-insensitive), if
@@ -516,19 +501,12 @@ class SanitizeKeyPoints(Transform):
 
     def __init__(
         self,
-        min_valid_edge_distance: int = 0,
-        min_invalid_points: Union[int, float] = 1,
         labels_getter: Union[Callable[[Any], Any], str, None] = None,
     ) -> None:
         super().__init__()
-        self.min_valid_edge_distance = min_valid_edge_distance
-        self.min_invalid_points = min_invalid_points
         self.labels_getter = labels_getter
         self._labels_getter = _parse_labels_getter(labels_getter)
 
-        if min_invalid_points <= 0:
-            raise ValueError(f"min_invalid_points must be > 0. Got {min_invalid_points}.")
-
     def forward(self, *inputs: Any) -> Any:
         inputs = inputs if len(inputs) > 1 else inputs[0]
 
@@ -559,8 +537,6 @@ def forward(self, *inputs: Any) -> Any:
         valid = F._misc._get_sanitize_keypoints_mask(
             points,
             canvas_size=points.canvas_size,
-            min_valid_edge_distance=self.min_valid_edge_distance,
-            min_invalid_points=self.min_invalid_points,
         )
 
         params = dict(valid=valid, labels=labels)