From 85ba8be9dd2adcdfd3daf7172453ac40470ee66c Mon Sep 17 00:00:00 2001 From: Copilot <198982749+Copilot@users.noreply.github.com> Date: Tue, 31 Mar 2026 01:48:57 +0200 Subject: [PATCH] Fix KeyPoints 2D boolean mask indexing (uniform-count selection) (#2188) * Fix KeyPoints 2D boolean mask filtering (keypoints[keypoints.confidence > 0.5]) * Fix ruff E501 and mypy type annotation in KeyPoints 2D mask branch * Add match= to pytest.raises to fix ruff PT011 * Add shape validation for 2D boolean mask in KeyPoints.__getitem__ * Add edge-case tests for KeyPoints 2D boolean mask filtering * refactor: extract 2D bool mask handling into _get_by_2d_bool_mask private method * Document 2D mask uniform-count requirement; add canonical single-object test --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: Borda <6035284+Borda@users.noreply.github.com> Co-authored-by: Claude Code Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- src/supervision/key_points/core.py | 76 +++++++++++++++++++ tests/key_points/test_core.py | 115 +++++++++++++++++++++++++++++ 2 files changed, 191 insertions(+) diff --git a/src/supervision/key_points/core.py b/src/supervision/key_points/core.py index a1636aa5..cd8190ae 100644 --- a/src/supervision/key_points/core.py +++ b/src/supervision/key_points/core.py @@ -670,6 +670,79 @@ class KeyPoints: else: return cls.empty() + def _get_by_2d_bool_mask(self, mask: npt.NDArray[np.bool_]) -> KeyPoints: + """Filter keypoints using a 2D boolean mask of shape `(n, m)`. + + This method selects the **same set of keypoints from every object**, so + every row of `mask` must contain the same number of `True` values. The + result is a new `KeyPoints` whose keypoint count is that uniform `k`. + + This is suitable for use cases such as *"keep only the left-side joints for + all persons"* — where the selected joint indices are identical across objects. + + It is **not** suitable for per-object confidence filtering + (`kp[kp.confidence > 0.5]`) when the threshold yields a different number of + passing keypoints per object, because NumPy cannot represent a ragged + `(n, ?, 2)` array. For that pattern either process objects individually or + zero out low-confidence entries in-place via `kp.confidence`. + + For the single-object case (`n == 1`) any boolean mask always satisfies the + uniform-count requirement, so `kp[kp.confidence > 0.5]` works as expected. + + Args: + mask: A boolean array of shape `(n, m)` where `n` is the number of + objects and `m` is the number of keypoints per object. Every row + must select the same number of keypoints so that the result can be + stored in a uniform `(n, k, ...)` array. + + Returns: + A new `KeyPoints` instance containing only the keypoints selected by + the mask for each object. + + Raises: + ValueError: If `mask.shape[0]` does not match the number of objects, if + `mask.shape[1]` does not match the number of keypoints, or if + different rows of the mask select different numbers of `True` values. + """ + n = len(self.xy) + if mask.shape[0] != n: + raise ValueError( + f"2D boolean mask row count {mask.shape[0]} does not match " + f"object count {n}." + ) + if mask.shape[1] != self.xy.shape[1]: + raise ValueError( + f"2D boolean mask column count {mask.shape[1]} does not match " + f"keypoint count {self.xy.shape[1]}." + ) + counts = np.sum(mask, axis=1) + if n > 0 and not np.all(counts == counts[0]): + raise ValueError( + "Cannot filter keypoints with a 2D boolean mask where rows have " + "different numbers of True values. " + "All objects must select the same number of keypoints. " + f"Got counts per object: {counts.tolist()}" + ) + k = int(counts[0]) if n > 0 else 0 + xy_selected = np.zeros((n, k, self.xy.shape[2]), dtype=self.xy.dtype) + conf_selected: npt.NDArray[np.float32] | None = None + if self.confidence is not None: + conf_selected = cast( + npt.NDArray[np.float32], + np.zeros((n, k), dtype=self.confidence.dtype), + ) + for row in range(n): + row_indices = np.flatnonzero(mask[row]) + xy_selected[row] = self.xy[row, row_indices] + if conf_selected is not None and self.confidence is not None: + conf_selected[row] = self.confidence[row, row_indices] + return KeyPoints( + xy=xy_selected, + confidence=conf_selected, + class_id=self.class_id.copy() if self.class_id is not None else None, + data=get_data_item(self.data, slice(None)), + ) + def __getitem__( self, index: Index1D | Index2D | str, @@ -677,6 +750,9 @@ class KeyPoints: if isinstance(index, str): return self.data.get(index) + if isinstance(index, np.ndarray) and index.ndim == 2 and index.dtype == bool: + return self._get_by_2d_bool_mask(index) + if not isinstance(index, tuple): index = (index, slice(None)) diff --git a/tests/key_points/test_core.py b/tests/key_points/test_core.py index 03feb571..bf10b1c2 100644 --- a/tests/key_points/test_core.py +++ b/tests/key_points/test_core.py @@ -251,6 +251,121 @@ KEY_POINTS = _create_key_points( _create_key_points(xy=[[[8, 9]]], confidence=[[0.5]], class_id=[0]), DoesNotRaise(), ), # select the last anchor from the first skeleton by index + ( + KEY_POINTS, + np.array( + [ + [True, False, True, False, False], + [True, True, False, False, False], + [False, True, True, False, False], + ] + ), + _create_key_points( + xy=[ + [[0, 1], [4, 5]], + [[10, 11], [12, 13]], + [[22, 23], [24, 25]], + ], + confidence=[[0.8, 0.6], [0.7, 0.9], [0.6, 0.8]], + class_id=[0, 1, 2], + ), + DoesNotRaise(), + ), # filter keypoints by 2D boolean mask, same count per row + ( + _create_key_points( + xy=[[[0, 1], [2, 3], [4, 5]]], + confidence=[[0.8, 0.2, 0.6]], + class_id=[0], + ), + np.array([[True, False, True]]), + _create_key_points( + xy=[[[0, 1], [4, 5]]], + confidence=[[0.8, 0.6]], + class_id=[0], + ), + DoesNotRaise(), + ), # filter keypoints by 2D boolean mask, single object + ( + _create_key_points( + xy=[ + [[0, 1], [2, 3], [4, 5]], + [[10, 11], [12, 13], [14, 15]], + ], + confidence=[ + [0.8, 0.2, 0.6], + [0.1, 0.2, 0.3], + ], + class_id=[0, 1], + ), + np.array([[True, False, True], [False, False, False]]), + None, + pytest.raises(ValueError, match="different numbers of True values"), + ), # 2D boolean mask with different counts per row raises ValueError + ( + _create_key_points( + xy=[[[0, 1], [2, 3], [4, 5]]], + class_id=[0], + ), + np.array([[True, False, True]]), + _create_key_points( + xy=[[[0, 1], [4, 5]]], + class_id=[0], + ), + DoesNotRaise(), + ), # 2D boolean mask with confidence=None — no confidence array in result + ( + _create_key_points( + xy=[[[0, 1], [2, 3], [4, 5]]], + confidence=[[0.8, 0.2, 0.6]], + class_id=[0], + ), + np.array([[True, False]]), + None, + pytest.raises(ValueError, match="column count"), + ), # 2D boolean mask column count mismatch raises ValueError + ( + _create_key_points( + xy=[[[0, 1], [2, 3], [4, 5]]], + confidence=[[0.8, 0.2, 0.6]], + class_id=[0], + ), + np.array([[True, False, True], [True, False, True]]), + None, + pytest.raises(ValueError, match="row count"), + ), # 2D boolean mask row count mismatch raises ValueError + ( + _create_key_points( + xy=[[[0, 1], [2, 3]], [[4, 5], [6, 7]]], + confidence=[[0.8, 0.2], [0.6, 0.9]], + class_id=[0, 1], + ), + np.array([[False, False], [False, False]]), + KeyPoints( + xy=np.zeros((2, 0, 2), dtype=np.float32), + confidence=np.zeros((2, 0), dtype=np.float32), + class_id=np.array([0, 1]), + ), + DoesNotRaise(), + ), # all-False 2D mask — all rows select 0 keypoints, equal counts → ok + ( + _create_key_points( + xy=[[[0, 1], [2, 3], [4, 5]]], + confidence=[[0.8, 0.2, 0.6]], + class_id=[0], + ), + _create_key_points( + xy=[[[0, 1], [2, 3], [4, 5]]], + confidence=[[0.8, 0.2, 0.6]], + class_id=[0], + ).confidence + > 0.5, + _create_key_points( + xy=[[[0, 1], [4, 5]]], + confidence=[[0.8, 0.6]], + class_id=[0], + ), + DoesNotRaise(), + ), # kp[kp.confidence > 0.5] — single-object canonical use case ], ) def test_key_points_getitem(key_points, index, expected_result, exception):