From 9ed1f07ca0c4db29b79a019ef3ff5d5363aa1021 Mon Sep 17 00:00:00 2001 From: Vijay Vignesh Date: Mon, 30 Mar 2026 14:36:31 -0400 Subject: [PATCH] Add unittests for `supervision.key_points.core` (#2190) * adding unit testcases for keypoints from_inference(), from_mediapipe() and from_yolo_nas() functions * modifying test_from_mediapipe_input() * Apply suggestions from code review --------- Co-authored-by: Jirka Borovec <6035284+Borda@users.noreply.github.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- tests/helpers.py | 43 ++++++++++++ tests/key_points/test_core.py | 129 +++++++++++++++++++++++++++++++++- 2 files changed, 171 insertions(+), 1 deletion(-) diff --git a/tests/helpers.py b/tests/helpers.py index 9ea1ef6e..1ce209b3 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -305,6 +305,49 @@ class _FakeYoloNasResults: self.prediction = prediction +class _FakeYoloNasKeyPoint: + """YOLO-NAS-like key point struct.""" + + def __init__(self, poses, labels=None): + self.poses = np.array(poses, dtype=np.float32) + if labels is not None: + self.labels = np.array(labels, dtype=int) + + +class _FakeYoloNasKeyPointResults: + """YOLO-NAS-like results exposing key points.""" + + def __init__(self, prediction: _FakeYoloNasKeyPoint, class_names=None): + self.prediction = prediction + self.class_names = class_names + + +class _FakeMediapipeLandmark: + def __init__(self, x, y, visibility=1.0): + self.x = x + self.y = y + self.visibility = visibility + + +class _FakeMediapipePose: + def __init__(self, landmarks: list[_FakeMediapipeLandmark]): + self.landmark = landmarks + + +class _FakeMediapipeResults: + def __init__( + self, + pose_landmarks: list[list[_FakeMediapipeLandmark]] + | _FakeMediapipePose + | None = None, + face_landmarks: _FakeMediapipeLandmark | None = None, + multi_face_landmarks: list[_FakeMediapipeLandmark] | None = None, + ): + self.pose_landmarks = pose_landmarks + self.face_landmarks = face_landmarks + self.multi_face_landmarks = multi_face_landmarks + + def create_yolo_dataset( dataset_dir: str, num_images: int = 15, diff --git a/tests/key_points/test_core.py b/tests/key_points/test_core.py index 59454d55..03feb571 100644 --- a/tests/key_points/test_core.py +++ b/tests/key_points/test_core.py @@ -4,7 +4,14 @@ import numpy as np import pytest from supervision.key_points.core import KeyPoints -from tests.helpers import _create_key_points +from tests.helpers import ( + _create_key_points, + _FakeMediapipeLandmark, + _FakeMediapipePose, + _FakeMediapipeResults, + _FakeYoloNasKeyPoint, + _FakeYoloNasKeyPointResults, +) KEY_POINTS = _create_key_points( xy=[ @@ -431,3 +438,123 @@ def test_key_points_equality_with_data(): ) 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