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 <noreply@anthropic.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
Copilot
GitHub
parent
129118817a
commit
85ba8be9dd
|
|
@ -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))
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Reference in New Issue