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

@@ -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,

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