From 5ced8c4ea0d3a4efee38ac7d2dc080f269a81c20 Mon Sep 17 00:00:00 2001 From: Jirka Borovec <6035284+Borda@users.noreply.github.com> Date: Fri, 30 Jan 2026 02:22:19 +0900 Subject: [PATCH] Add unit tests for `key_points` - skeleton & annotators (#2099) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Refactor key point annotator tests to use reusable fixtures for test clarity and maintainability * Apply suggestions from code review * fix(pre_commit): 🎨 auto format pre-commit hooks * Move `assert_image_mostly_same` helper to a shared module and update imports in tests * Update test/key_points/test_skeletons.py --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- test/annotators/test_core.py | 32 +----------- test/conftest.py | 66 ++++++++++++++++++++++++ test/helpers.py | 30 +++++++++++ test/key_points/test_annotators.py | 81 ++++++++++++++++++++++++++++++ test/key_points/test_skeletons.py | 60 ++++++++++++++++++++++ 5 files changed, 238 insertions(+), 31 deletions(-) create mode 100644 test/conftest.py create mode 100644 test/key_points/test_annotators.py create mode 100644 test/key_points/test_skeletons.py diff --git a/test/annotators/test_core.py b/test/annotators/test_core.py index 636a2236..0b07830d 100644 --- a/test/annotators/test_core.py +++ b/test/annotators/test_core.py @@ -31,7 +31,7 @@ from supervision.annotators.utils import ColorLookup from supervision.detection.core import Detections from supervision.draw.color import Color from supervision.geometry.core import Position -from test.helpers import _create_detections +from test.helpers import _create_detections, assert_image_mostly_same @pytest.fixture @@ -58,36 +58,6 @@ def gradient_image() -> np.ndarray: return image -def assert_image_mostly_same( - original: np.ndarray, annotated: np.ndarray, similarity_threshold: float = 0.9 -) -> None: - """ - Assert that the annotated image is mostly the same as the original. - - Args: - original: Original image - annotated: Annotated image - similarity_threshold: - Minimum percentage of pixels that should be the same (0.0 to 1.0) - """ - # Check that images have the same shape - assert original.shape == annotated.shape - - # Calculate number of identical pixels - identical_pixels = np.sum(np.all(original == annotated, axis=-1)) - total_pixels = original.shape[0] * original.shape[1] - similarity = identical_pixels / total_pixels - - # Check that at least similarity_threshold of pixels are identical - assert similarity >= similarity_threshold, ( - f"Images are only {similarity:.1%} similar, " - f"which is below the {similarity_threshold:.1%} threshold" - ) - - # Check that the image is not completely identical - assert not np.array_equal(original, annotated), "Images are completely identical" - - class TestBoxAnnotator: """Tests for BoxAnnotator class""" diff --git a/test/conftest.py b/test/conftest.py new file mode 100644 index 00000000..cd03031e --- /dev/null +++ b/test/conftest.py @@ -0,0 +1,66 @@ +import numpy as np +import pytest + +import supervision as sv +from test.helpers import _create_key_points + + +@pytest.fixture +def scene() -> np.ndarray: + return np.zeros((100, 100, 3), dtype=np.uint8) + + +@pytest.fixture +def sample_key_points() -> sv.KeyPoints: + return _create_key_points( + xy=[ + [ + [10, 10], + [20, 20], + [30, 30], + [40, 40], + [50, 50], + [60, 60], + [70, 70], + [80, 80], + [90, 90], + [10, 20], + [20, 30], + [30, 40], + [40, 50], + [50, 60], + [60, 70], + [70, 80], + [80, 90], + ], + [ + [10, 40], + [20, 50], + [30, 60], + [40, 70], + [50, 80], + [60, 90], + [70, 10], + [80, 20], + [90, 30], + [10, 50], + [20, 60], + [30, 70], + [40, 80], + [50, 90], + [60, 10], + [70, 20], + [80, 30], + ], + ], + confidence=[ + [0.8] * 17, + [0.6] * 17, + ], + class_id=[0, 1], + ) + + +@pytest.fixture +def empty_key_points() -> sv.KeyPoints: + return sv.KeyPoints.empty() diff --git a/test/helpers.py b/test/helpers.py index 11fb1266..bff7c867 100644 --- a/test/helpers.py +++ b/test/helpers.py @@ -208,3 +208,33 @@ def assert_almost_equal(actual, expected, tolerance=1e-5): AssertionError: Expected 0.5, but got 0.6. """ assert abs(actual - expected) < tolerance, f"Expected {expected}, but got {actual}." + + +def assert_image_mostly_same( + original: np.ndarray, annotated: np.ndarray, similarity_threshold: float = 0.9 +) -> None: + """ + Assert that the annotated image is mostly the same as the original. + + Args: + original: Original image + annotated: Annotated image + similarity_threshold: + Minimum percentage of pixels that should be the same (0.0 to 1.0) + """ + # Check that images have the same shape + assert original.shape == annotated.shape + + # Calculate number of identical pixels + identical_pixels = np.sum(np.all(original == annotated, axis=-1)) + total_pixels = original.shape[0] * original.shape[1] + similarity = identical_pixels / total_pixels + + # Check that at least similarity_threshold of pixels are identical + assert similarity >= similarity_threshold, ( + f"Images are only {similarity:.1%} similar, " + f"which is below the {similarity_threshold:.1%} threshold" + ) + + # Check that the image is not completely identical + assert not np.array_equal(original, annotated), "Images are completely identical" diff --git a/test/key_points/test_annotators.py b/test/key_points/test_annotators.py new file mode 100644 index 00000000..db2f8413 --- /dev/null +++ b/test/key_points/test_annotators.py @@ -0,0 +1,81 @@ +import numpy as np + +import supervision as sv +from test.helpers import assert_image_mostly_same + + +class TestVertexAnnotator: + def test_annotate_with_default_parameters(self, scene, sample_key_points): + """Test annotation with default parameters.""" + annotator = sv.VertexAnnotator() + result = annotator.annotate(scene=scene.copy(), key_points=sample_key_points) + + # Check that the scene has been modified + assert_image_mostly_same( + original=scene, annotated=result, similarity_threshold=0.8 + ) + + def test_annotate_with_custom_color_and_radius(self, scene, sample_key_points): + """Test annotation with custom color and radius.""" + color = sv.Color.RED + radius = 5 + annotator = sv.VertexAnnotator(color=color, radius=radius) + result = annotator.annotate(scene=scene.copy(), key_points=sample_key_points) + + # Check that the scene has been modified + assert_image_mostly_same( + original=scene, annotated=result, similarity_threshold=0.7 + ) + + def test_annotate_empty_key_points(self, scene, empty_key_points): + """Test annotation with empty key points returns unchanged scene.""" + annotator = sv.VertexAnnotator() + result = annotator.annotate(scene=scene.copy(), key_points=empty_key_points) + + # Should return the original scene unchanged + assert np.array_equal(result, scene) + + +class TestEdgeAnnotator: + def test_annotate_with_default_parameters(self, scene, sample_key_points): + """Test annotation with default parameters using COCO skeleton.""" + annotator = sv.EdgeAnnotator() + result = annotator.annotate(scene=scene.copy(), key_points=sample_key_points) + + # Check that the scene has been modified + assert_image_mostly_same( + original=scene, annotated=result, similarity_threshold=0.7 + ) + + def test_annotate_with_custom_edges(self, scene, sample_key_points): + """Test annotation with custom edge definitions.""" + edges = [(1, 2), (2, 3)] + annotator = sv.EdgeAnnotator(edges=edges) + result = annotator.annotate(scene=scene.copy(), key_points=sample_key_points) + + # Check that the scene has been modified + assert_image_mostly_same( + original=scene, annotated=result, similarity_threshold=0.8 + ) + + def test_annotate_empty_key_points(self, scene, empty_key_points): + """Test annotation with empty key points returns unchanged scene.""" + annotator = sv.EdgeAnnotator() + result = annotator.annotate(scene=scene.copy(), key_points=empty_key_points) + + # Should return the original scene unchanged + assert np.array_equal(result, scene) + + def test_annotate_no_edges_found(self, scene): + """Test annotation when no matching skeleton is found.""" + # Key points with more vertices than any skeleton + large_key_points = sv.KeyPoints( + xy=np.array([[[i * 10, i * 10] for i in range(100)]], dtype=np.float32), + confidence=np.array([[0.8] * 100], dtype=np.float32), + class_id=np.array([0], dtype=int), + ) + annotator = sv.EdgeAnnotator() + result = annotator.annotate(scene=scene.copy(), key_points=large_key_points) + + # Should return the original scene unchanged (no edges found) + assert np.array_equal(result, scene) diff --git a/test/key_points/test_skeletons.py b/test/key_points/test_skeletons.py new file mode 100644 index 00000000..e9756288 --- /dev/null +++ b/test/key_points/test_skeletons.py @@ -0,0 +1,60 @@ +from supervision.key_points.skeletons import ( + SKELETONS_BY_EDGE_COUNT, + SKELETONS_BY_VERTEX_COUNT, + Skeleton, +) + + +class TestSkeletons: + def test_skeleton_enum_values(self): + """Test skeleton enum has correct structure.""" + for skeleton in Skeleton: + assert isinstance(skeleton.value, tuple) + assert all( + isinstance(edge, tuple) and len(edge) == 2 for edge in skeleton.value + ) + + def test_skeletons_by_vertex_count(self): + """Test SKELETONS_BY_VERTEX_COUNT dictionary population.""" + # Test that the dictionary is populated + assert len(SKELETONS_BY_VERTEX_COUNT) > 0 + + # Test specific known skeletons + coco_skeleton = Skeleton.COCO.value + assert 17 in SKELETONS_BY_VERTEX_COUNT # COCO has 17 keypoints + assert SKELETONS_BY_VERTEX_COUNT[17] == coco_skeleton + + def test_skeletons_by_edge_count(self): + """Test SKELETONS_BY_EDGE_COUNT dictionary mapping.""" + # Test that the dictionary is populated + assert len(SKELETONS_BY_EDGE_COUNT) > 0 + + # Reconstruct the expected mapping: for each skeleton, map the number of + # edges in skeleton.value to skeleton.value itself (as done in skeletons.py). + expected = {} + for skeleton in Skeleton: + edge_count = len(skeleton.value) + expected[edge_count] = skeleton.value + + assert SKELETONS_BY_EDGE_COUNT == expected + + def test_unique_vertices_calculation(self): + """Test unique vertices calculation from skeleton edges.""" + coco_skeleton = Skeleton.COCO.value + unique_vertices = {vertex for edge in coco_skeleton for vertex in edge} + assert len(unique_vertices) == 17 # COCO has 17 keypoints + + def test_skeletons_by_vertex_count_mapping_behaviour(self): + """Test SKELETONS_BY_VERTEX_COUNT uses last-in-wins for duplicate counts.""" + expected_mapping = {} + for skeleton in Skeleton: + vertex_count = len({v for edge in skeleton.value for v in edge}) + # Mimic skeletons.py: later skeletons overwrite earlier ones + expected_mapping[vertex_count] = skeleton.value + + # The keys (vertex counts) should match + assert set(SKELETONS_BY_VERTEX_COUNT.keys()) == set(expected_mapping.keys()) + + # For each vertex count, the stored skeleton should be the last one encountered + for vertex_count, skeleton_value in expected_mapping.items(): + assert SKELETONS_BY_VERTEX_COUNT[vertex_count] == skeleton_value