feat: support per-class skeleton edges in EdgeAnnotator via dict[int,…#2293
Merged
SkalskiP merged 15 commits intoJun 4, 2026
Conversation
Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
There was a problem hiding this comment.
Pull request overview
This PR extends Supervision’s keypoint annotation utilities to better support datasets that contain multiple skeleton types by allowing EdgeAnnotator to select skeleton edges per detection class, while keeping existing single-skeleton and auto-detect behaviors for backward compatibility.
Changes:
- Expanded
EdgeAnnotator.edgesto acceptdict[int, Sequence[tuple[int, int]]]keyed byclass_id, and adjusted failure handling to skip only the failing instance rather than aborting the whole annotation pass. - Refactored
VertexLabelAnnotatorto support per-classlabelsviadict[int, list[str]]and to skip vertices usingkey_points.visiblewhen provided. - Minor cleanup in tests and repo config (comment removals in tests; ignore a local script).
Reviewed changes
Copilot reviewed 2 out of 3 changed files in this pull request and generated 6 comments.
| File | Description |
|---|---|
src/supervision/key_points/annotators.py |
Adds per-class skeleton edge selection in EdgeAnnotator; refactors/extends VertexLabelAnnotator (per-class labels + visibility handling); docstring example blocks removed. |
tests/key_points/test_annotators.py |
Small test comment cleanup; no new coverage added for the new per-class behaviors. |
.gitignore |
Ignores scripts/test_multi_skeleton.py. |
Comment on lines
+549
to
+551
| if isinstance(colors, list): | ||
| return colors | ||
| return [colors] * points_count |
Comment on lines
63
to
66
| Returns: | ||
| The annotated image, matching the type of `scene` (`numpy.ndarray` | ||
| or `PIL.Image.Image`) | ||
|
|
||
| Example: | ||
| ```pycon | ||
| >>> import numpy as np | ||
| >>> import supervision as sv | ||
| >>> image = np.zeros((100, 100, 3), dtype=np.uint8) | ||
| >>> key_points = sv.KeyPoints( | ||
| ... xy=np.array([[[50, 50], [60, 60]]], dtype=np.float32) | ||
| ... ) | ||
| >>> vertex_annotator = sv.VertexAnnotator( | ||
| ... color=sv.Color.GREEN, | ||
| ... radius=10 | ||
| ... ) | ||
| >>> annotated_frame = vertex_annotator.annotate( | ||
| ... scene=image.copy(), | ||
| ... key_points=key_points | ||
| ... ) | ||
| >>> annotated_frame.shape | ||
| (100, 100, 3) | ||
|
|
||
| ``` | ||
| """ |
Comment on lines
+148
to
+152
| if class_id is None or class_id not in self.edges: | ||
| logger.warning("No edges defined for class_id=%s", class_id) | ||
| continue | ||
| edges = self.edges[class_id] | ||
| elif self.edges: |
Comment on lines
397
to
401
| self, | ||
| scene: ImageType, | ||
| key_points: KeyPoints, | ||
| labels: list[str] | None = None, | ||
| labels: list[str] | dict[int, list[str]] | None = None, | ||
| ) -> ImageType: |
Comment on lines
397
to
401
| self, | ||
| scene: ImageType, | ||
| key_points: KeyPoints, | ||
| labels: list[str] | None = None, | ||
| labels: list[str] | dict[int, list[str]] | None = None, | ||
| ) -> ImageType: |
Collaborator
Author
There was a problem hiding this comment.
Just updated the PR description. It reflects both EdgeAnnotator and VertexLabelAnnotator changes now.
Borda
reviewed
Jun 4, 2026
Borda
approved these changes
Jun 4, 2026
…n syntax highlighting and multi-skeleton demos
…or _resolve_labels
Comment on lines
+213
to
+222
| if isinstance(self.edges, dict): | ||
| class_id = ( | ||
| key_points.class_id[detection_index] | ||
| if key_points.class_id is not None | ||
| else None | ||
| ) | ||
| if class_id is None or class_id not in self.edges: | ||
| logger.warning("No edges defined for class_id=%s", class_id) | ||
| continue | ||
| edges = self.edges[class_id] |
Comment on lines
+597
to
+600
| class_id = ( | ||
| key_points.class_id[i] if key_points.class_id is not None else None | ||
| ) | ||
| instance_labels = self._resolve_labels(labels, points_count, class_id) |
Comment on lines
+213
to
+230
| if isinstance(self.edges, dict): | ||
| class_id = ( | ||
| key_points.class_id[detection_index] | ||
| if key_points.class_id is not None | ||
| else None | ||
| ) | ||
| if class_id is None or class_id not in self.edges: | ||
| logger.warning("No edges defined for class_id=%s", class_id) | ||
| continue | ||
| edges = self.edges[class_id] | ||
| elif self.edges: | ||
| edges = self.edges | ||
| else: | ||
| _looked_up = SKELETONS_BY_VERTEX_COUNT.get(len(xy)) | ||
| if not _looked_up: | ||
| logger.warning("No skeleton found with %d vertices", len(xy)) | ||
| continue | ||
| edges = _looked_up |
Comment on lines
495
to
+514
| def annotate( | ||
| self, | ||
| scene: ImageType, | ||
| key_points: KeyPoints, | ||
| labels: list[str] | None = None, | ||
| labels: list[str] | dict[int, list[str]] | None = None, | ||
| ) -> ImageType: | ||
| """ | ||
| A class that draws labels of skeleton vertices on images. It uses specified key | ||
| points to determine the locations where the vertices should be drawn. | ||
| Draws labels at skeleton vertex positions on the image. Vertices | ||
| marked not visible via ``key_points.visible`` are skipped. | ||
|
|
||
| Args: | ||
| scene: The image where vertex labels will be drawn. `ImageType` is a | ||
| flexible type, accepting either `numpy.ndarray` or `PIL.Image.Image`. | ||
| key_points: A collection of key points where each key point consists of x | ||
| and y coordinates. | ||
| labels: A list of labels to be displayed on the annotated image. If not | ||
| provided, keypoint indices will be used. | ||
| labels: Labels to display at each keypoint. If ``None``, keypoint | ||
| indices are used. A ``list[str]`` applies the same labels to | ||
| every instance. A ``dict[int, list[str]]`` maps ``class_id`` | ||
| to per-class label lists, enabling correct labeling for | ||
| datasets with multiple skeleton types. |
1eea095
into
refactor/keypoints-visibility-confidence-split
6 checks passed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
EdgeAnnotator.edgesnow accepts adict[int, Sequence[tuple[int, int]]]mappingclass_idto skeleton edges, enabling correct rendering for datasets with multiple skeleton types (e.g. people and animals with different keypoint counts).Sequence) and auto-detect (None) modes remain unchanged for backward compatibility.return scenechanged tocontinue).labelsparameter now accepts adict[int, list[str]]mappingclass_idto per-class label lists.key_points.visibleinstead of relying on theanchors != 0heuristic, which could skip legitimate keypoints at the origin or label padded keypoints at non-zero coordinates.