Files
wehub-resource-sync 9194ef5abd
Docs/Test Workflow / Test docs build (push) Failing after 0s
Check links & references / links-check (push) Failing after 1s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.10) (push) Failing after 0s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.11) (push) Failing after 0s
PR Conflict Labeler / main (push) Failing after 2s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.12) (push) Failing after 2s
Pytest/Test Workflow / Import Test and Pytest Run (ubuntu-latest, 3.13) (push) Failing after 0s
Pytest/Test Workflow / Build this Package (push) Failing after 5s
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.10) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.11) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.12) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (macos-latest, 3.13) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.10) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.11) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.12) (push) Has been cancelled
Pytest/Test Workflow / Import Test and Pytest Run (windows-latest, 3.13) (push) Has been cancelled
Pytest/Test Workflow / testing-guardian (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:06:10 +08:00

1813 lines
59 KiB
Python

from contextlib import nullcontext as DoesNotRaise
import numpy as np
import pytest
from supervision.key_points.core import KeyPoints
from tests.helpers import (
_create_key_points,
_FakeMediapipeLandmark,
_FakeMediapipePose,
_FakeMediapipeResults,
_FakeYoloNasKeyPoint,
_FakeYoloNasKeyPointResults,
)
KEY_POINTS = _create_key_points(
xy=[
[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]],
[[10, 11], [12, 13], [14, 15], [16, 17], [18, 19]],
[[20, 21], [22, 23], [24, 25], [26, 27], [28, 29]],
],
confidence=[
[0.8, 0.2, 0.6, 0.1, 0.5],
[0.7, 0.9, 0.3, 0.4, 0.0],
[0.1, 0.6, 0.8, 0.2, 0.7],
],
class_id=[0, 1, 2],
)
@pytest.mark.parametrize(
("key_points", "index", "expected_result", "exception"),
[
(
KeyPoints.empty(),
slice(None),
KeyPoints.empty(),
DoesNotRaise(),
), # slice all key points when key points object empty
(
KEY_POINTS,
slice(None),
KEY_POINTS,
DoesNotRaise(),
), # slice all key points when key points object nonempty
(
KEY_POINTS,
slice(0, 1),
_create_key_points(
xy=[[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]],
confidence=[[0.8, 0.2, 0.6, 0.1, 0.5]],
class_id=[0],
),
DoesNotRaise(),
), # select the first skeleton by slice
(
KEY_POINTS,
slice(0, 2),
_create_key_points(
xy=[
[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]],
[[10, 11], [12, 13], [14, 15], [16, 17], [18, 19]],
],
confidence=[
[0.8, 0.2, 0.6, 0.1, 0.5],
[0.7, 0.9, 0.3, 0.4, 0.0],
],
class_id=[0, 1],
),
DoesNotRaise(),
), # select the first skeleton by slice
(
KEY_POINTS,
0,
_create_key_points(
xy=[[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]],
confidence=[[0.8, 0.2, 0.6, 0.1, 0.5]],
class_id=[0],
),
DoesNotRaise(),
), # select the first skeleton by index
(
KEY_POINTS,
-1,
_create_key_points(
xy=[[[20, 21], [22, 23], [24, 25], [26, 27], [28, 29]]],
confidence=[[0.1, 0.6, 0.8, 0.2, 0.7]],
class_id=[2],
),
DoesNotRaise(),
), # select the last skeleton by index
(
KEY_POINTS,
[0, 1],
_create_key_points(
xy=[
[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]],
[[10, 11], [12, 13], [14, 15], [16, 17], [18, 19]],
],
confidence=[
[0.8, 0.2, 0.6, 0.1, 0.5],
[0.7, 0.9, 0.3, 0.4, 0.0],
],
class_id=[0, 1],
),
DoesNotRaise(),
), # select the first two skeletons by index; list
(
KEY_POINTS,
np.array([0, 1]),
_create_key_points(
xy=[
[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]],
[[10, 11], [12, 13], [14, 15], [16, 17], [18, 19]],
],
confidence=[
[0.8, 0.2, 0.6, 0.1, 0.5],
[0.7, 0.9, 0.3, 0.4, 0.0],
],
class_id=[0, 1],
),
DoesNotRaise(),
), # select the first two skeletons by index; np.array
(
KEY_POINTS,
[True, True, False],
_create_key_points(
xy=[
[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]],
[[10, 11], [12, 13], [14, 15], [16, 17], [18, 19]],
],
confidence=[
[0.8, 0.2, 0.6, 0.1, 0.5],
[0.7, 0.9, 0.3, 0.4, 0.0],
],
class_id=[0, 1],
),
DoesNotRaise(),
), # select only skeletons associated with positive filter; list
(
KEY_POINTS,
np.array([True, True, False]),
_create_key_points(
xy=[
[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]],
[[10, 11], [12, 13], [14, 15], [16, 17], [18, 19]],
],
confidence=[
[0.8, 0.2, 0.6, 0.1, 0.5],
[0.7, 0.9, 0.3, 0.4, 0.0],
],
class_id=[0, 1],
),
DoesNotRaise(),
), # select only skeletons associated with positive filter; list
(
KEY_POINTS,
(slice(None), slice(None)),
KEY_POINTS,
DoesNotRaise(),
), # slice all anchors from all skeletons
(
KEY_POINTS,
(slice(None), slice(0, 1)),
_create_key_points(
xy=[[[0, 1]], [[10, 11]], [[20, 21]]],
confidence=[[0.8], [0.7], [0.1]],
class_id=[0, 1, 2],
),
DoesNotRaise(),
), # slice the first anchor from every skeleton
(
KEY_POINTS,
(slice(None), slice(0, 2)),
_create_key_points(
xy=[[[0, 1], [2, 3]], [[10, 11], [12, 13]], [[20, 21], [22, 23]]],
confidence=[[0.8, 0.2], [0.7, 0.9], [0.1, 0.6]],
class_id=[0, 1, 2],
),
DoesNotRaise(),
), # slice the first anchor two anchors from every skeleton
(
KEY_POINTS,
(slice(None), 0),
_create_key_points(
xy=[[[0, 1]], [[10, 11]], [[20, 21]]],
confidence=[[0.8], [0.7], [0.1]],
class_id=[0, 1, 2],
),
DoesNotRaise(),
), # select the first anchor from every skeleton by index
(
KEY_POINTS,
(slice(None), -1),
_create_key_points(
xy=[[[8, 9]], [[18, 19]], [[28, 29]]],
confidence=[[0.5], [0.0], [0.7]],
class_id=[0, 1, 2],
),
DoesNotRaise(),
), # select the last anchor from every skeleton by index
(
KEY_POINTS,
(slice(None), [0, 1]),
_create_key_points(
xy=[[[0, 1], [2, 3]], [[10, 11], [12, 13]], [[20, 21], [22, 23]]],
confidence=[[0.8, 0.2], [0.7, 0.9], [0.1, 0.6]],
class_id=[0, 1, 2],
),
DoesNotRaise(),
), # select the first two anchors from every skeleton by index; list
(
KEY_POINTS,
(slice(None), np.array([0, 1])),
_create_key_points(
xy=[[[0, 1], [2, 3]], [[10, 11], [12, 13]], [[20, 21], [22, 23]]],
confidence=[[0.8, 0.2], [0.7, 0.9], [0.1, 0.6]],
class_id=[0, 1, 2],
),
DoesNotRaise(),
), # select the first two anchors from every skeleton by index; np.array
(
KEY_POINTS,
(slice(None), [True, True, False, False, False]),
_create_key_points(
xy=[[[0, 1], [2, 3]], [[10, 11], [12, 13]], [[20, 21], [22, 23]]],
confidence=[[0.8, 0.2], [0.7, 0.9], [0.1, 0.6]],
class_id=[0, 1, 2],
),
DoesNotRaise(),
), # select only anchors associated with positive filter; list
(
KEY_POINTS,
(slice(None), np.array([True, True, False, False, False])),
_create_key_points(
xy=[[[0, 1], [2, 3]], [[10, 11], [12, 13]], [[20, 21], [22, 23]]],
confidence=[[0.8, 0.2], [0.7, 0.9], [0.1, 0.6]],
class_id=[0, 1, 2],
),
DoesNotRaise(),
), # select only anchors associated with positive filter; np.array
(
KEY_POINTS,
(0, 0),
_create_key_points(xy=[[[0, 1]]], confidence=[[0.8]], class_id=[0]),
DoesNotRaise(),
), # select the first anchor from the first skeleton by index
(
KEY_POINTS,
(0, -1),
_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),
keypoint_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],
).keypoint_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):
with exception:
result = key_points[index]
assert result == expected_result
def test_key_points_select_returns_subset() -> None:
"""Select returns a typed KeyPoints subset for row indexes."""
result = KEY_POINTS.select([0, 2])
assert result == _create_key_points(
xy=[
[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]],
[[20, 21], [22, 23], [24, 25], [26, 27], [28, 29]],
],
confidence=[
[0.8, 0.2, 0.6, 0.1, 0.5],
[0.1, 0.6, 0.8, 0.2, 0.7],
],
class_id=[0, 2],
)
def test_key_points_get_data_returns_data_value() -> None:
"""Get data returns the stored data value or None."""
key_points = _create_key_points(
xy=[[[0, 1], [2, 3]]],
confidence=[[0.8, 0.2]],
class_id=[0],
)
key_points["custom_data"] = ["value1"]
assert key_points.get_data("custom_data") == ["value1"]
assert key_points.get_data("missing") is None
KEY_POINTS_WITH_DET_CONF = _create_key_points(
xy=[
[[0, 1], [2, 3], [4, 5]],
[[10, 11], [12, 13], [14, 15]],
[[20, 21], [22, 23], [24, 25]],
],
confidence=[
[0.8, 0.2, 0.6],
[0.7, 0.9, 0.3],
[0.1, 0.6, 0.8],
],
class_id=[0, 1, 0],
detection_confidence=[0.95, 0.40, 0.85],
)
@pytest.mark.parametrize(
("key_points", "index", "expected_result"),
[
pytest.param(
KEY_POINTS_WITH_DET_CONF,
KEY_POINTS_WITH_DET_CONF.detection_confidence > 0.5,
_create_key_points(
xy=[
[[0, 1], [2, 3], [4, 5]],
[[20, 21], [22, 23], [24, 25]],
],
confidence=[[0.8, 0.2, 0.6], [0.1, 0.6, 0.8]],
class_id=[0, 0],
detection_confidence=[0.95, 0.85],
),
id="filter-by-detection-confidence-threshold",
),
pytest.param(
KEY_POINTS_WITH_DET_CONF,
KEY_POINTS_WITH_DET_CONF.class_id == 0,
_create_key_points(
xy=[
[[0, 1], [2, 3], [4, 5]],
[[20, 21], [22, 23], [24, 25]],
],
confidence=[[0.8, 0.2, 0.6], [0.1, 0.6, 0.8]],
class_id=[0, 0],
detection_confidence=[0.95, 0.85],
),
id="filter-by-class-id",
),
pytest.param(
KEY_POINTS_WITH_DET_CONF,
KEY_POINTS_WITH_DET_CONF.class_id == 1,
_create_key_points(
xy=[[[10, 11], [12, 13], [14, 15]]],
confidence=[[0.7, 0.9, 0.3]],
class_id=[1],
detection_confidence=[0.40],
),
id="filter-by-class-id-single-result",
),
pytest.param(
KEY_POINTS_WITH_DET_CONF,
(KEY_POINTS_WITH_DET_CONF.detection_confidence > 0.5)
& (KEY_POINTS_WITH_DET_CONF.class_id == 0),
_create_key_points(
xy=[
[[0, 1], [2, 3], [4, 5]],
[[20, 21], [22, 23], [24, 25]],
],
confidence=[[0.8, 0.2, 0.6], [0.1, 0.6, 0.8]],
class_id=[0, 0],
detection_confidence=[0.95, 0.85],
),
id="filter-by-detection-confidence-and-class-id",
),
pytest.param(
KEY_POINTS_WITH_DET_CONF,
KEY_POINTS_WITH_DET_CONF.detection_confidence > 0.99,
KeyPoints(
xy=np.zeros((0, 3, 2), dtype=np.float32),
keypoint_confidence=np.zeros((0, 3), dtype=np.float32),
detection_confidence=np.array([], dtype=np.float32),
class_id=np.array([], dtype=int),
),
id="filter-all-out-returns-empty",
),
pytest.param(
KEY_POINTS_WITH_DET_CONF,
KEY_POINTS_WITH_DET_CONF.class_id == 99,
KeyPoints(
xy=np.zeros((0, 3, 2), dtype=np.float32),
keypoint_confidence=np.zeros((0, 3), dtype=np.float32),
detection_confidence=np.array([], dtype=np.float32),
class_id=np.array([], dtype=int),
),
id="filter-by-nonexistent-class-returns-empty",
),
pytest.param(
KeyPoints.empty(),
np.array([], dtype=bool),
KeyPoints.empty(),
id="filter-empty-keypoints-stays-empty",
),
pytest.param(
KEY_POINTS_WITH_DET_CONF,
KEY_POINTS_WITH_DET_CONF.detection_confidence > 0.0,
KEY_POINTS_WITH_DET_CONF,
id="filter-keeps-all-when-all-pass",
),
pytest.param(
KEY_POINTS_WITH_DET_CONF,
np.int64(0),
_create_key_points(
xy=[[[0, 1], [2, 3], [4, 5]]],
confidence=[[0.8, 0.2, 0.6]],
class_id=[0],
detection_confidence=[0.95],
),
id="np-integer-scalar-with-det-conf",
),
pytest.param(
KEY_POINTS_WITH_DET_CONF,
np.array(0),
_create_key_points(
xy=[[[0, 1], [2, 3], [4, 5]]],
confidence=[[0.8, 0.2, 0.6]],
class_id=[0],
detection_confidence=[0.95],
),
id="0d-ndarray-with-det-conf",
),
],
)
def test_key_points_getitem_detection_level(key_points, index, expected_result):
"""Detection-level filtering mirrors Detections API patterns."""
result = key_points[index]
assert result == expected_result
class TestKeyPointsVisible:
"""Tests for the `visible` mask field on KeyPoints."""
def test_visible_defaults_to_none(self):
kp = _create_key_points(
xy=[[[0, 1], [2, 3]]],
confidence=[[0.9, 0.8]],
class_id=[0],
)
assert kp.visible is None
def test_visible_set_from_confidence_threshold(self):
kp = _create_key_points(
xy=[[[0, 1], [2, 3], [4, 5]]],
confidence=[[0.9, 0.1, 0.6]],
class_id=[0],
)
kp.visible = kp.keypoint_confidence > 0.5
expected = np.array([[True, False, True]])
np.testing.assert_array_equal(kp.visible, expected)
def test_visible_preserved_on_skeleton_filter(self):
kp = _create_key_points(
xy=[
[[0, 1], [2, 3]],
[[10, 11], [12, 13]],
],
confidence=[[0.9, 0.1], [0.3, 0.8]],
class_id=[0, 1],
detection_confidence=[0.95, 0.40],
visible=[[True, False], [False, True]],
)
result = kp[kp.detection_confidence > 0.5]
assert result.visible is not None
np.testing.assert_array_equal(result.visible, np.array([[True, False]]))
def test_visible_preserved_on_int_index(self):
kp = _create_key_points(
xy=[[[0, 1], [2, 3]], [[10, 11], [12, 13]]],
confidence=[[0.9, 0.1], [0.3, 0.8]],
class_id=[0, 1],
visible=[[True, False], [False, True]],
)
result = kp[0]
assert result.visible is not None
np.testing.assert_array_equal(result.visible, np.array([[True, False]]))
def test_visible_preserved_on_anchor_slice(self):
kp = _create_key_points(
xy=[[[0, 1], [2, 3], [4, 5]]],
confidence=[[0.9, 0.1, 0.6]],
class_id=[0],
visible=[[True, False, True]],
)
result = kp[:, [0, 2]]
assert result.visible is not None
np.testing.assert_array_equal(result.visible, np.array([[True, True]]))
def test_visible_preserved_on_2d_bool_mask(self):
kp = _create_key_points(
xy=[
[[0, 1], [2, 3], [4, 5]],
[[10, 11], [12, 13], [14, 15]],
],
confidence=[[0.9, 0.1, 0.6], [0.7, 0.2, 0.8]],
class_id=[0, 1],
visible=[[True, False, True], [True, False, True]],
)
mask = np.array([[True, False, True], [True, False, True]])
result = kp[mask]
assert result.visible is not None
np.testing.assert_array_equal(
result.visible, np.array([[True, True], [True, True]])
)
def test_visible_none_stays_none_on_filter(self):
kp = _create_key_points(
xy=[[[0, 1], [2, 3]], [[10, 11], [12, 13]]],
class_id=[0, 1],
)
result = kp[0]
assert result.visible is None
def test_equality_with_visible(self):
kp1 = _create_key_points(
xy=[[[0, 1], [2, 3]]],
class_id=[0],
visible=[[True, False]],
)
kp2 = _create_key_points(
xy=[[[0, 1], [2, 3]]],
class_id=[0],
visible=[[True, False]],
)
kp3 = _create_key_points(
xy=[[[0, 1], [2, 3]]],
class_id=[0],
visible=[[False, True]],
)
assert kp1 == kp2
assert kp1 != kp3
def test_equality_visible_none_vs_set(self):
kp1 = _create_key_points(
xy=[[[0, 1], [2, 3]]],
class_id=[0],
)
kp2 = _create_key_points(
xy=[[[0, 1], [2, 3]]],
class_id=[0],
visible=[[True, True]],
)
assert kp1 != kp2
def test_key_points_empty():
"""Test the creation and behavior of an empty KeyPoints object."""
empty_key_points = KeyPoints.empty()
assert len(empty_key_points) == 0
assert empty_key_points.is_empty()
assert empty_key_points.xy.shape == (0, 0, 2)
def test_key_points_is_empty():
"""Test the is_empty method for KeyPoints objects."""
empty_key_points = KeyPoints.empty()
assert empty_key_points.is_empty()
non_empty_key_points = _create_key_points(
xy=[[[0, 1], [2, 3]]],
confidence=[[0.8, 0.9]],
class_id=[0],
)
assert not non_empty_key_points.is_empty()
def test_key_points_zero_length_slice_is_empty():
"""A filtered KeyPoints object with zero rows is empty."""
key_points = _create_key_points(
xy=[[[0, 1], [2, 3]]],
confidence=[[0.8, 0.9]],
class_id=[0],
)
assert key_points[np.array([], dtype=int)].is_empty()
def test_key_points_setitem():
"""Test the __setitem__ method for KeyPoints objects."""
key_points = _create_key_points(
xy=[[[0, 1], [2, 3]]],
confidence=[[0.8, 0.9]],
class_id=[0],
)
key_points["custom_data"] = ["value1"]
assert "custom_data" in key_points.data
assert np.array_equal(key_points.data["custom_data"], np.array(["value1"]))
with pytest.raises(TypeError, match=r"Value must be a np\.ndarray or a list"):
key_points["invalid_data"] = 123
@pytest.mark.parametrize(
("key_points", "expected_xyxy", "expected_confidence", "expected_class_id"),
[
(
_create_key_points(
xy=[[[0, 1], [2, 3], [4, 5]]],
confidence=[[0.8, 0.9, 0.7]],
class_id=[0],
),
np.array([[0, 1, 4, 5]], dtype=np.float32),
np.array([0.8], dtype=np.float32),
np.array([0]),
),
(
_create_key_points(
xy=[[[0, 0], [2, 3], [4, 5]]],
confidence=[[0.8, 0.9, 0.7]],
class_id=[0],
),
np.array([[2, 3, 4, 5]], dtype=np.float32),
np.array([0.8], dtype=np.float32),
np.array([0]),
),
],
)
def test_key_points_as_detections(
key_points, expected_xyxy, expected_confidence, expected_class_id
):
"""Test the as_detections method for KeyPoints objects."""
detections = key_points.as_detections()
assert len(detections) == len(expected_xyxy)
assert np.array_equal(detections.xyxy, expected_xyxy)
assert np.allclose(detections.confidence, expected_confidence)
assert np.array_equal(detections.class_id, expected_class_id)
def test_key_points_as_detections_empty():
"""Test the as_detections method for empty KeyPoints objects."""
empty_key_points = KeyPoints.empty()
empty_detections = empty_key_points.as_detections()
assert empty_detections.is_empty()
def test_key_points_as_detections_ignores_missing_keypoints():
"""A [0, 0] keypoint is treated as missing and excluded from the box."""
key_points = _create_key_points(
xy=[[[0, 0], [10, 20], [30, 40]]],
confidence=[[0.0, 0.8, 0.6]],
class_id=[0],
)
detections = key_points.as_detections()
assert np.array_equal(detections.xyxy, np.array([[10, 20, 30, 40]]))
def test_key_points_as_detections_uses_detection_confidence():
"""detection_confidence is preferred over the keypoint-confidence mean."""
key_points = _create_key_points(
xy=[[[10, 20], [30, 40]]],
confidence=[[0.1, 0.2]],
class_id=[0],
detection_confidence=[0.95],
)
detections = key_points.as_detections()
assert np.allclose(detections.confidence, np.array([0.95], dtype=np.float32))
def test_key_points_as_detections_selected_keypoint_indices():
"""Only the selected keypoints contribute to the bounding box."""
key_points = _create_key_points(
xy=[[[0, 0], [10, 20], [30, 40], [100, 100]]],
confidence=[[0.5, 0.8, 0.6, 0.9]],
class_id=[0],
)
detections = key_points.as_detections(selected_keypoint_indices=[1, 2])
assert np.array_equal(detections.xyxy, np.array([[10, 20, 30, 40]]))
def test_key_points_as_detections_confidence_over_selected_indices():
"""Confidence mean uses only the selected keypoint columns, not all."""
key_points = _create_key_points(
xy=[[[0, 0], [10, 20], [30, 40], [100, 100]]],
confidence=[[0.5, 0.8, 0.6, 0.9]],
class_id=[0],
)
detections = key_points.as_detections(selected_keypoint_indices=[1, 2])
expected_confidence = np.mean([0.8, 0.6], dtype=np.float32)
assert np.isclose(detections.confidence[0], expected_confidence)
def test_key_points_as_detections_mixed_valid_invalid_batch():
"""Batch with one all-zero skeleton: invalid skeleton gets box zeroed."""
key_points = _create_key_points(
xy=[[[0, 0], [0, 0]], [[10, 20], [30, 40]]],
confidence=[[0.0, 0.0], [0.8, 0.6]],
class_id=[0, 1],
)
detections = key_points.as_detections()
# Only the skeleton with at least one valid keypoint survives
assert len(detections) == 1
assert np.array_equal(detections.xyxy, np.array([[10, 20, 30, 40]]))
def test_key_points_getitem_empty_list():
"""Selecting with an empty list returns an empty KeyPoints, like Detections."""
key_points = _create_key_points(
xy=[[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
class_id=[0, 1],
)
result = key_points[[]]
assert len(result) == 0
assert result.is_empty()
@pytest.mark.parametrize(
"selected_keypoint_indices",
[
pytest.param(np.array([0, 1]), id="numpy-array"),
pytest.param(iter([0, 1]), id="generator"),
pytest.param((0, 1), id="tuple"),
],
)
def test_key_points_as_detections_index_container_types(selected_keypoint_indices):
"""selected_keypoint_indices accepts any Iterable[int] container type."""
key_points = _create_key_points(
xy=[[[10, 10], [20, 20], [30, 15]]],
class_id=[0],
)
detections = key_points.as_detections(
selected_keypoint_indices=selected_keypoint_indices,
)
assert np.array_equal(detections.xyxy, np.array([[10, 10, 20, 20]]))
def test_key_points_as_detections_empty_indices_selects_all():
"""An empty selected_keypoint_indices behaves like None (selects all)."""
key_points = _create_key_points(
xy=[[[10, 10], [20, 20], [30, 15]]],
confidence=[[0.1, 0.3, 0.5]],
class_id=[0],
detection_confidence=[0.9],
)
key_points["custom_data"] = ["person"]
all_selected = key_points.as_detections()
empty_list_selected = key_points.as_detections(selected_keypoint_indices=[])
assert np.array_equal(all_selected.xyxy, empty_list_selected.xyxy)
assert np.array_equal(all_selected.confidence, empty_list_selected.confidence)
assert np.array_equal(all_selected.class_id, empty_list_selected.class_id)
assert np.array_equal(
all_selected.data["custom_data"], empty_list_selected.data["custom_data"]
)
@pytest.mark.parametrize(
("xy", "expected_xyxy"),
[
pytest.param(
[[[10, 20], [30, 20]]],
np.array([[10, 20, 30, 20]]),
id="collinear-keypoints",
),
pytest.param(
[[[15, 25]]],
np.array([[15, 25, 15, 25]]),
id="single-keypoint",
),
],
)
def test_key_points_as_detections_keeps_degenerate_skeletons(xy, expected_xyxy):
"""A skeleton with valid keypoints keeps its box even when the area is zero."""
key_points = _create_key_points(xy=xy, class_id=[0])
detections = key_points.as_detections()
assert len(detections) == 1
assert np.array_equal(detections.xyxy, expected_xyxy)
def test_key_points_as_detections_filters_invalid_and_aligns_degenerate_metadata():
"""Invalid skeletons are removed while valid degenerate metadata stays aligned."""
key_points = _create_key_points(
xy=[
[[0, 0], [0, 0]],
[[np.nan, 5], [0, 0]],
[[np.inf, 5], [0, 0]],
[[15, 25], [np.nan, 50]],
[[10, 20], [30, 20]],
],
class_id=[0, 1, 2, 3, 4],
detection_confidence=[0.1, 0.2, 0.3, 0.4, 0.5],
)
key_points["custom_data"] = ["zero", "nan", "inf", "point", "line"]
detections = key_points.as_detections()
assert np.array_equal(
detections.xyxy,
np.array([[15, 25, 15, 25], [10, 20, 30, 20]], dtype=np.float32),
)
assert np.array_equal(detections.class_id, np.array([3, 4]))
assert np.array_equal(detections.confidence, np.array([0.4, 0.5], dtype=np.float32))
assert np.array_equal(detections.data["custom_data"], np.array(["point", "line"]))
def test_key_points_as_detections_with_data():
"""Test the as_detections method preserves data."""
key_points = _create_key_points(
xy=[[[0, 1], [2, 3], [4, 5]]],
confidence=[[0.8, 0.9, 0.7]],
class_id=[0],
)
key_points["custom_data"] = ["value1"]
detections = key_points.as_detections()
assert "custom_data" in detections.data
assert np.array_equal(detections.data["custom_data"], np.array(["value1"]))
@pytest.mark.parametrize(
("key_points_list", "expected_result", "exception"),
[
(
[],
KeyPoints.empty(),
DoesNotRaise(),
), # empty list
(
[KeyPoints.empty(), KeyPoints.empty()],
KeyPoints.empty(),
DoesNotRaise(),
), # only empty KeyPoints
(
[
_create_key_points(
xy=[[[10, 10], [20, 20]]],
confidence=[[0.9, 0.8]],
class_id=[0],
),
],
_create_key_points(
xy=[[[10, 10], [20, 20]]],
confidence=[[0.9, 0.8]],
class_id=[0],
),
DoesNotRaise(),
), # single KeyPoints
(
[
KeyPoints.empty(),
_create_key_points(
xy=[[[10, 10], [20, 20]]],
confidence=[[0.9, 0.8]],
class_id=[0],
),
KeyPoints.empty(),
],
_create_key_points(
xy=[[[10, 10], [20, 20]]],
confidence=[[0.9, 0.8]],
class_id=[0],
),
DoesNotRaise(),
), # empty KeyPoints are ignored
(
[
KeyPoints(xy=np.array([[[10, 10], [20, 20]]], dtype=np.float32)),
KeyPoints(xy=np.array([[[30, 30], [40, 40]]], dtype=np.float32)),
],
KeyPoints(
xy=np.array(
[[[10, 10], [20, 20]], [[30, 30], [40, 40]]], dtype=np.float32
)
),
DoesNotRaise(),
), # xy only; all optional fields None
(
[
_create_key_points(
xy=[[[10, 10], [20, 20]]],
confidence=[[0.9, 0.8]],
class_id=[0],
detection_confidence=[0.95],
visible=[[True, False]],
data={"class_name": ["person"]},
),
_create_key_points(
xy=[[[30, 30], [40, 40]], [[50, 50], [60, 60]]],
confidence=[[0.7, 0.6], [0.5, 0.4]],
class_id=[1, 2],
detection_confidence=[0.85, 0.75],
visible=[[True, True], [False, True]],
data={"class_name": ["dog", "cat"]},
),
],
_create_key_points(
xy=[
[[10, 10], [20, 20]],
[[30, 30], [40, 40]],
[[50, 50], [60, 60]],
],
confidence=[[0.9, 0.8], [0.7, 0.6], [0.5, 0.4]],
class_id=[0, 1, 2],
detection_confidence=[0.95, 0.85, 0.75],
visible=[[True, False], [True, True], [False, True]],
data={"class_name": ["person", "dog", "cat"]},
),
DoesNotRaise(),
), # all fields populated on every input
(
[
_create_key_points(
xy=[[[10, 10], [20, 20]]],
confidence=[[0.9, 0.8]],
class_id=[0],
),
KeyPoints(xy=np.array([[[30, 30], [40, 40]]], dtype=np.float32)),
],
None,
pytest.raises(ValueError, match="All or none of the 'class_id'"),
), # class_id set on one input and None on the other
(
[
_create_key_points(
xy=[[[10, 10], [20, 20]]],
visible=[[True, False]],
),
KeyPoints(xy=np.array([[[30, 30], [40, 40]]], dtype=np.float32)),
],
None,
pytest.raises(ValueError, match="All or none of the 'visible'"),
), # visible set on one input and None on the other
(
[
_create_key_points(
xy=[[[10, 10], [20, 20]]],
detection_confidence=[0.9],
),
KeyPoints(xy=np.array([[[30, 30], [40, 40]]], dtype=np.float32)),
],
None,
pytest.raises(
ValueError, match="All or none of the 'detection_confidence'"
),
), # detection_confidence set on one input and None on the other
(
[
KeyPoints(xy=np.array([[[10, 10], [20, 20]]], dtype=np.float32)),
KeyPoints(
xy=np.array([[[10, 10], [20, 20], [30, 30]]], dtype=np.float32)
),
],
None,
pytest.raises(ValueError, match="same number of keypoints"),
), # mismatched keypoint counts per skeleton
(
[
KeyPoints(xy=np.array([[[10, 10], [20, 20]]], dtype=np.float32)),
KeyPoints(
xy=np.array([[[30, 30, 0.9], [40, 40, 0.8]]], dtype=np.float32)
),
],
None,
pytest.raises(ValueError, match=r"got depths \[2, 3\]"),
), # mismatched coordinate depths per skeleton
(
[
_create_key_points(
xy=[[[10, 10], [20, 20]]],
data={"class_name": ["person"]},
),
_create_key_points(
xy=[[[30, 30], [40, 40]]],
data={"tracker_id": [7]},
),
],
None,
pytest.raises(ValueError, match="same keys to merge"),
), # data dictionaries with mismatched keys
pytest.param(
[
KeyPoints(xy=np.zeros((2, 0, 2), dtype=np.float32)),
KeyPoints(xy=np.zeros((3, 0, 2), dtype=np.float32)),
],
KeyPoints(xy=np.zeros((5, 0, 2), dtype=np.float32)),
DoesNotRaise(),
id="zero-keypoints-merge-succeeds",
),
pytest.param(
[
KeyPoints(xy=np.zeros((2, 0, 2), dtype=np.float32)),
KeyPoints(xy=np.zeros((1, 2, 2), dtype=np.float32)),
],
None,
pytest.raises(ValueError, match="same number of keypoints"),
id="zero-vs-nonzero-keypoints-mismatch",
),
],
)
def test_key_points_merge(
key_points_list: list[KeyPoints],
expected_result: KeyPoints | None,
exception: Exception,
) -> None:
"""Test KeyPoints.merge field merging, mismatches, and zero-keypoint cases."""
with exception:
result = KeyPoints.merge(key_points_list=key_points_list)
assert result == expected_result, f"Expected: {expected_result}, Got: {result}"
def test_key_points_iteration():
"""Test the iteration over KeyPoints objects."""
key_points = _create_key_points(
xy=[[[0, 1], [2, 3]], [[4, 5], [6, 7]]],
confidence=[[0.8, 0.9], [0.7, 0.6]],
class_id=[0, 1],
)
iterations = 0
for i, (xy, kp_confidence, class_id, data) in enumerate(key_points):
iterations += 1
assert xy.shape == (2, 2)
assert kp_confidence.shape == (2,)
assert class_id in [0, 1]
assert isinstance(data, dict)
assert iterations == 2
def test_key_points_iteration_no_confidence():
"""Test the iteration over KeyPoints objects without confidence."""
key_points_no_conf = _create_key_points(
xy=[[[0, 1], [2, 3]]],
confidence=None,
class_id=[0],
)
for xy, kp_confidence, class_id, data in key_points_no_conf:
assert kp_confidence is None
def test_key_points_merge_then_with_nms_deduplicates_overlapping_detections():
"""Test merge-then-NMS removes duplicated overlapping skeletons."""
key_points_list = [
_create_key_points(
xy=[[[100, 100], [200, 200]], [[400, 400], [500, 500]]],
detection_confidence=[0.9, 0.6],
class_id=[0, 0],
),
_create_key_points(
xy=[[[100, 100], [200, 200]], [[700, 700], [800, 800]]],
detection_confidence=[0.85, 0.7],
class_id=[0, 0],
),
]
merged = KeyPoints.merge(key_points_list)
result = merged.with_nms(threshold=0.5, class_agnostic=False)
assert len(merged) == 4
assert len(result) == 3
assert len(result) < len(merged)
np.testing.assert_allclose(
np.sort(result.detection_confidence),
np.array([0.6, 0.7, 0.9], dtype=np.float32),
)
@pytest.mark.parametrize(
("key_points1", "key_points2", "expected_equal"),
[
(
_create_key_points(
xy=[[[0, 1], [2, 3]]], confidence=[[0.8, 0.9]], class_id=[0]
),
_create_key_points(
xy=[[[0, 1], [2, 3]]], confidence=[[0.8, 0.9]], class_id=[0]
),
True,
),
(
_create_key_points(
xy=[[[0, 1], [2, 3]]], confidence=[[0.8, 0.9]], class_id=[0]
),
_create_key_points(
xy=[[[0, 1], [2, 3]]], confidence=[[0.8, 0.9]], class_id=[1]
),
False,
),
(
_create_key_points(
xy=[[[0, 1], [2, 3]]], confidence=[[0.8, 0.9]], class_id=[0]
),
_create_key_points(
xy=[[[0, 1], [2, 4]]], confidence=[[0.8, 0.9]], class_id=[0]
),
False,
),
(
_create_key_points(
xy=[[[0, 1], [2, 3]]], confidence=[[0.8, 0.9]], class_id=[0]
),
_create_key_points(
xy=[[[0, 1], [2, 3]]], confidence=[[0.8, 0.8]], class_id=[0]
),
False,
),
],
)
def test_key_points_equality(key_points1, key_points2, expected_equal):
"""Test the equality comparison for KeyPoints objects."""
status = key_points1 == key_points2
assert status is expected_equal
def test_key_points_equality_with_data():
"""Test the equality comparison for KeyPoints objects with data."""
key_points1 = _create_key_points(
xy=[[[0, 1], [2, 3]]], confidence=[[0.8, 0.9]], class_id=[0]
)
key_points2 = _create_key_points(
xy=[[[0, 1], [2, 3]]], confidence=[[0.8, 0.9]], class_id=[0]
)
key_points2["custom"] = ["value"]
assert key_points1 != key_points2
@pytest.mark.parametrize(
("inference_results", "expected_key_points"),
[
(
{
"predictions": [
{
"class_id": 1,
"class": "person",
"keypoints": [
{"x": 100, "y": 150, "confidence": 0.9},
{"x": 120, "y": 160, "confidence": 0.85},
],
}
]
},
_create_key_points(
xy=[[[100.0, 150.0], [120.0, 160.0]]],
confidence=[[0.9, 0.85]],
class_id=[1],
data={"class_name": np.array(["person"])},
),
),
({"predictions": []}, KeyPoints.empty()),
],
)
def test_from_inference_input(inference_results, expected_key_points):
"""Test the from_inference method with valid input."""
key_points = KeyPoints.from_inference(inference_results)
assert key_points == expected_key_points
def test_from_inference_invalid_input():
"""Test the from_inference method with invalid input."""
key_points = _create_key_points(
xy=[[[0, 1], [2, 3]]], confidence=[[0.8, 0.9]], class_id=[0]
)
with pytest.raises(
ValueError, match=r"from_inference\(\) operates on a single result at a time.*"
):
KeyPoints.from_inference([key_points])
@pytest.mark.parametrize(
("yolo_nas_results", "expected_key_points"),
[
(
_FakeYoloNasKeyPointResults(
_FakeYoloNasKeyPoint(
poses=[[[100.0, 150.0, 0.9], [120.0, 160.0, 0.85]]],
labels=[1],
),
),
_create_key_points(
xy=[[[100.0, 150.0], [120.0, 160.0]]],
confidence=[[0.9, 0.85]],
class_id=[1],
),
),
(
_FakeYoloNasKeyPointResults(
_FakeYoloNasKeyPoint(
poses=[],
),
),
KeyPoints.empty(),
),
],
)
def test_from_yolo_nas_input(yolo_nas_results, expected_key_points):
"""Test the from_yolo_nas method with valid input."""
key_points = KeyPoints.from_yolo_nas(yolo_nas_results)
assert key_points == expected_key_points
@pytest.mark.parametrize(
("mediapipe_results", "resolution_wh", "expected_key_points"),
[
(
_FakeMediapipeResults(
pose_landmarks=_FakeMediapipePose(
landmarks=[
_FakeMediapipeLandmark(0.5, 0.75, 0.9),
_FakeMediapipeLandmark(0.6, 0.8, 0.85),
]
)
),
(200, 200),
_create_key_points(
xy=[[[100.0, 150.0], [120.0, 160.0]]],
confidence=[[0.9, 0.85]],
class_id=None,
),
),
(
_FakeMediapipeResults(
pose_landmarks=[
[
_FakeMediapipeLandmark(0.5, 0.75, 0.9),
_FakeMediapipeLandmark(0.6, 0.8, 0.85),
]
]
),
(200, 200),
_create_key_points(
xy=[[[100.0, 150.0], [120.0, 160.0]]],
confidence=[[0.9, 0.85]],
class_id=None,
),
),
],
)
def test_from_mediapipe_input(mediapipe_results, resolution_wh, expected_key_points):
"""Test the from_mediapipe method with valid input."""
key_points = KeyPoints.from_mediapipe(
mediapipe_results, resolution_wh=resolution_wh
)
assert key_points == expected_key_points
def test_from_mediapipe_unknown_result_type_raises() -> None:
"""Unsupported Mediapipe result objects raise a descriptive ValueError."""
with pytest.raises(ValueError, match="Unsupported MediaPipe result"):
KeyPoints.from_mediapipe(object(), resolution_wh=(100, 100))
class TestDeprecatedConfidenceConstructor:
"""Tests for backward-compatible `confidence=` kwarg in KeyPoints()."""
def test_constructor_accepts_and_warns_on_deprecated_confidence_kwarg(self):
"""Deprecated confidence= warns and maps value to keypoint_confidence."""
from supervision.utils.internal import SupervisionWarnings
xy = np.array([[[1.0, 2.0], [3.0, 4.0]]], dtype=np.float32)
confidence = np.array([[0.9, 0.8]], dtype=np.float32)
with pytest.warns(SupervisionWarnings, match="deprecated since"):
key_points = KeyPoints(xy=xy, confidence=confidence)
np.testing.assert_array_equal(key_points.keypoint_confidence, confidence)
assert key_points.xy is xy
@pytest.mark.parametrize(
"kwargs",
[
pytest.param(
{"confidence": None, "keypoint_confidence": None},
id="confidence-first",
),
pytest.param(
{"keypoint_confidence": None, "confidence": None},
id="keypoint-confidence-first",
),
],
)
def test_constructor_rejects_both_confidence_and_keypoint_confidence(
self, kwargs: dict
):
"""ValueError raised regardless of kwarg order when both are passed."""
xy = np.array([[[1.0, 2.0], [3.0, 4.0]]], dtype=np.float32)
confidence_arr = np.array([[0.9, 0.8]], dtype=np.float32)
actual_kwargs = {k: confidence_arr for k in kwargs}
with pytest.raises(ValueError, match="Cannot pass both"):
KeyPoints(xy=xy, **actual_kwargs)
def test_constructor_normal_keypoint_confidence_path(self):
"""Normal keypoint_confidence= path works and emits no deprecation warning."""
import warnings
xy = np.array([[[1.0, 2.0], [3.0, 4.0]]], dtype=np.float32)
kp_conf = np.array([[0.9, 0.8]], dtype=np.float32)
with warnings.catch_warnings():
warnings.simplefilter("error")
kp = KeyPoints(xy=xy, keypoint_confidence=kp_conf)
np.testing.assert_array_equal(kp.keypoint_confidence, kp_conf)
assert kp.data == {}
def test_constructor_confidence_none_does_not_warn(self):
"""Explicit confidence=None is silently ignored — no warning emitted."""
import warnings
xy = np.array([[[1.0, 2.0], [3.0, 4.0]]], dtype=np.float32)
with warnings.catch_warnings():
warnings.simplefilter("error")
kp = KeyPoints(xy=xy, confidence=None)
assert kp.keypoint_confidence is None
def test_constructor_data_none_defaults_to_empty_dict(self):
"""Explicit data=None normalizes to empty dict, not None."""
xy = np.array([[[1.0, 2.0]]], dtype=np.float32)
assert KeyPoints(xy=xy).data == {}
assert KeyPoints(xy=xy, data=None).data == {}
def test_keypoints_init_covers_all_dataclass_fields(self):
"""Custom __init__ must assign every dataclass field — guards against drift."""
import dataclasses
import inspect
field_names = {f.name for f in dataclasses.fields(KeyPoints)}
init_params = set(inspect.signature(KeyPoints.__init__).parameters) - {
"self",
"confidence",
}
assert field_names == init_params, (
f"Field/init drift: {field_names.symmetric_difference(init_params)}"
)
@pytest.mark.parametrize(
("key_points", "threshold", "class_agnostic", "expected_result"),
[
pytest.param(
KeyPoints.empty(),
0.5,
True,
KeyPoints.empty(),
id="empty",
),
pytest.param(
_create_key_points(
xy=[[[10, 20], [30, 40]]],
detection_confidence=[0.9],
class_id=[0],
),
0.5,
False,
_create_key_points(
xy=[[[10, 20], [30, 40]]],
detection_confidence=[0.9],
class_id=[0],
),
id="single-skeleton",
),
pytest.param(
_create_key_points(
xy=[
[[10, 10], [20, 20]],
[[500, 500], [600, 600]],
],
detection_confidence=[0.9, 0.8],
class_id=[0, 0],
),
0.5,
False,
_create_key_points(
xy=[
[[10, 10], [20, 20]],
[[500, 500], [600, 600]],
],
detection_confidence=[0.9, 0.8],
class_id=[0, 0],
),
id="two-non-overlapping",
),
pytest.param(
_create_key_points(
xy=[
[[100, 100], [200, 200]],
[[150, 150], [250, 250]],
],
detection_confidence=[0.9, 0.7],
class_id=[0, 0],
),
0.9,
False,
_create_key_points(
xy=[
[[100, 100], [200, 200]],
[[150, 150], [250, 250]],
],
detection_confidence=[0.9, 0.7],
class_id=[0, 0],
),
id="two-overlapping-below-threshold",
),
pytest.param(
_create_key_points(
xy=[
[[100, 100], [200, 200]],
[[110, 110], [210, 210]],
],
detection_confidence=[0.9, 0.7],
class_id=[0, 0],
),
0.3,
False,
_create_key_points(
xy=[[[100, 100], [200, 200]]],
detection_confidence=[0.9],
class_id=[0],
),
id="two-overlapping-above-threshold",
),
pytest.param(
_create_key_points(
xy=[
[[100, 100], [200, 200]],
[[110, 110], [210, 210]],
[[500, 500], [600, 600]],
],
detection_confidence=[0.9, 0.7, 0.8],
class_id=[0, 0, 0],
),
0.3,
False,
_create_key_points(
xy=[
[[100, 100], [200, 200]],
[[500, 500], [600, 600]],
],
detection_confidence=[0.9, 0.8],
class_id=[0, 0],
),
id="three-skeletons-two-overlap",
),
pytest.param(
_create_key_points(
xy=[
[[100, 100], [200, 200]],
[[110, 110], [210, 210]],
],
detection_confidence=[0.9, 0.7],
class_id=[0, 1],
),
0.3,
False,
_create_key_points(
xy=[
[[100, 100], [200, 200]],
[[110, 110], [210, 210]],
],
detection_confidence=[0.9, 0.7],
class_id=[0, 1],
),
id="class-aware-different-classes-kept",
),
pytest.param(
_create_key_points(
xy=[
[[100, 100], [200, 200]],
[[110, 110], [210, 210]],
],
detection_confidence=[0.9, 0.7],
class_id=[0, 1],
),
0.3,
True,
_create_key_points(
xy=[[[100, 100], [200, 200]]],
detection_confidence=[0.9],
class_id=[0],
),
id="class-agnostic-suppresses-across-classes",
),
pytest.param(
_create_key_points(
xy=[
[[0, 0], [100, 100], [200, 200]],
[[0, 0], [110, 110], [210, 210]],
],
detection_confidence=[0.9, 0.7],
class_id=[0, 0],
),
0.3,
False,
_create_key_points(
xy=[[[0, 0], [100, 100], [200, 200]]],
detection_confidence=[0.9],
class_id=[0],
),
id="zero-keypoints-excluded-from-bbox",
),
pytest.param(
_create_key_points(
xy=[
[[100, 100], [200, 200], [0, 0], [0, 0]],
[[0, 0], [0, 0], [110, 110], [210, 210]],
],
detection_confidence=[0.9, 0.7],
class_id=[0, 1],
),
0.3,
False,
_create_key_points(
xy=[
[[100, 100], [200, 200], [0, 0], [0, 0]],
[[0, 0], [0, 0], [110, 110], [210, 210]],
],
detection_confidence=[0.9, 0.7],
class_id=[0, 1],
),
id="multi-skeleton-schema-class-aware",
),
pytest.param(
_create_key_points(
xy=[
[[100, 100], [200, 200], [0, 0], [0, 0]],
[[0, 0], [0, 0], [110, 110], [210, 210]],
],
detection_confidence=[0.9, 0.7],
class_id=[0, 1],
),
0.3,
True,
_create_key_points(
xy=[[[100, 100], [200, 200], [0, 0], [0, 0]]],
detection_confidence=[0.9],
class_id=[0],
),
id="multi-skeleton-schema-class-agnostic",
),
pytest.param(
_create_key_points(
xy=[
[[0, 0], [0, 0]],
[[100, 100], [200, 200]],
],
detection_confidence=[0.9, 0.8],
class_id=[0, 0],
),
0.5,
False,
_create_key_points(
xy=[
[[0, 0], [0, 0]],
[[100, 100], [200, 200]],
],
detection_confidence=[0.9, 0.8],
class_id=[0, 0],
),
id="all-zero-skeleton-passes-through",
),
pytest.param(
_create_key_points(
xy=[
[[100, 100], [200, 200]],
[[110, 110], [210, 210]],
],
detection_confidence=[0.9, 0.7],
class_id=[0, 0],
visible=[[True, False], [False, True]],
),
0.3,
False,
_create_key_points(
xy=[
[[100, 100], [200, 200]],
[[110, 110], [210, 210]],
],
detection_confidence=[0.9, 0.7],
class_id=[0, 0],
visible=[[True, False], [False, True]],
),
id="visible-mask-excludes-keypoints-from-bbox",
),
pytest.param(
_create_key_points(
xy=[
[[100, 100], [0, 0]],
[[300, 300], [0, 0]],
],
detection_confidence=[0.9, 0.8],
class_id=[0, 0],
),
0.3,
False,
_create_key_points(
xy=[
[[100, 100], [0, 0]],
[[300, 300], [0, 0]],
],
detection_confidence=[0.9, 0.8],
class_id=[0, 0],
),
id="single-valid-keypoint-zero-area-bbox",
),
pytest.param(
_create_key_points(
xy=[
[[100, 100], [200, 200]],
[[110, 110], [210, 210]],
],
detection_confidence=[0.9, 0.7],
class_id=[0, 0],
),
0.0,
False,
_create_key_points(
xy=[[[100, 100], [200, 200]]],
detection_confidence=[0.9],
class_id=[0],
),
id="threshold-zero-suppresses-any-overlap",
),
pytest.param(
_create_key_points(
xy=[
[[100, 100], [200, 200]],
[[110, 110], [210, 210]],
],
detection_confidence=[0.9, 0.7],
class_id=[0, 0],
),
1.0,
False,
_create_key_points(
xy=[
[[100, 100], [200, 200]],
[[110, 110], [210, 210]],
],
detection_confidence=[0.9, 0.7],
class_id=[0, 0],
),
id="threshold-one-keeps-all",
),
],
)
def test_with_nms(key_points, threshold, class_agnostic, expected_result):
"""NMS filters overlapping keypoint skeletons."""
result = key_points.with_nms(threshold=threshold, class_agnostic=class_agnostic)
assert result == expected_result
@pytest.mark.parametrize(
("key_points", "threshold", "class_agnostic", "match"),
[
pytest.param(
_create_key_points(
xy=[[[10, 20], [30, 40]]],
class_id=[0],
),
0.5,
False,
"detection_confidence",
id="no-detection-confidence",
),
pytest.param(
_create_key_points(
xy=[[[10, 20], [30, 40]]],
detection_confidence=[0.9],
),
0.5,
False,
"class_id",
id="no-class-id-class-aware",
),
pytest.param(
_create_key_points(
xy=[[[10, 20], [30, 40]]],
class_id=[0],
),
0.5,
True,
"detection_confidence",
id="no-detection-confidence-class-agnostic",
),
],
)
def test_with_nms_raises(key_points, threshold, class_agnostic, match):
"""NMS raises when required fields are missing."""
with pytest.raises(ValueError, match=match):
key_points.with_nms(threshold=threshold, class_agnostic=class_agnostic)