From beaf89cb22f0aee82f71774fddb06beda7449293 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Mon, 28 Jul 2025 08:24:43 +0200 Subject: [PATCH] initial version of advanced key points filtering and slicing + some naming refactor --- docs/how_to/track_objects.md | 4 +- docs/keypoint/annotators.md | 12 +- docs/keypoint/core.md | 2 +- supervision/__init__.py | 4 +- .../{keypoint => key_points}/__init__.py | 0 .../{keypoint => key_points}/annotators.py | 4 +- supervision/{keypoint => key_points}/core.py | 208 +++++--- .../{keypoint => key_points}/skeletons.py | 0 supervision/validators/__init__.py | 6 +- test/key_points/__init__.py | 0 test/key_points/test_core.py | 444 ++++++++++++++++++ test/test_utils.py | 11 +- 12 files changed, 603 insertions(+), 92 deletions(-) rename supervision/{keypoint => key_points}/__init__.py (100%) rename supervision/{keypoint => key_points}/annotators.py (99%) rename supervision/{keypoint => key_points}/core.py (81%) rename supervision/{keypoint => key_points}/skeletons.py (100%) create mode 100644 test/key_points/__init__.py create mode 100644 test/key_points/test_core.py diff --git a/docs/how_to/track_objects.md b/docs/how_to/track_objects.md index c2be3cef..9bf17e86 100644 --- a/docs/how_to/track_objects.md +++ b/docs/how_to/track_objects.md @@ -345,7 +345,7 @@ Supervision is versatile and compatible with various models. Check this [link](/ We will define a `callback` function, which will process each frame of the video by obtaining model predictions and then annotating the frame based on these predictions. -Let's immediately visualize the results with our [`EdgeAnnotator`](/latest/keypoint/annotators/#supervision.keypoint.annotators.EdgeAnnotator) and [`VertexAnnotator`](https://supervision.roboflow.com/latest/keypoint/annotators/#supervision.keypoint.annotators.VertexAnnotator). +Let's immediately visualize the results with our [`EdgeAnnotator`](/latest/keypoint/annotators/#supervision.key_points.annotators.EdgeAnnotator) and [`VertexAnnotator`](https://supervision.roboflow.com/latest/keypoint/annotators/#supervision.key_points.annotators.VertexAnnotator). === "Ultralytics" @@ -408,7 +408,7 @@ Let's immediately visualize the results with our [`EdgeAnnotator`](/latest/keypo ### Convert to Detections -Keypoint tracking is currently supported via the conversion of `KeyPoints` to `Detections`. This is achieved with the [`KeyPoints.as_detections()`](/latest/keypoint/core/#supervision.keypoint.core.KeyPoints.as_detections) function. +Keypoint tracking is currently supported via the conversion of `KeyPoints` to `Detections`. This is achieved with the [`KeyPoints.as_detections()`](/latest/keypoint/core/#supervision.key_points.core.KeyPoints.as_detections) function. Let's convert to detections and visualize the results with our [`BoxAnnotator`](/latest/detection/annotators/#supervision.annotators.core.BoxAnnotator). diff --git a/docs/keypoint/annotators.md b/docs/keypoint/annotators.md index 32f30626..92c7ceba 100644 --- a/docs/keypoint/annotators.md +++ b/docs/keypoint/annotators.md @@ -78,19 +78,19 @@ comments: true
-

VertexAnnotator

+

VertexAnnotator

-:::supervision.keypoint.annotators.VertexAnnotator +:::supervision.key_points.annotators.VertexAnnotator
-

EdgeAnnotator

+

EdgeAnnotator

-:::supervision.keypoint.annotators.EdgeAnnotator +:::supervision.key_points.annotators.EdgeAnnotator
-

VertexLabelAnnotator

+

VertexLabelAnnotator

-:::supervision.keypoint.annotators.VertexLabelAnnotator +:::supervision.key_points.annotators.VertexLabelAnnotator diff --git a/docs/keypoint/core.md b/docs/keypoint/core.md index 7354baba..acb13e15 100644 --- a/docs/keypoint/core.md +++ b/docs/keypoint/core.md @@ -5,4 +5,4 @@ status: new # Keypoint Detection -:::supervision.keypoint.core.KeyPoints +:::supervision.key_points.core.KeyPoints diff --git a/supervision/__init__.py b/supervision/__init__.py index cc54fe70..ab45651a 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -108,12 +108,12 @@ from supervision.draw.utils import ( ) from supervision.geometry.core import Point, Position, Rect from supervision.geometry.utils import get_polygon_center -from supervision.keypoint.annotators import ( +from supervision.key_points.annotators import ( EdgeAnnotator, VertexAnnotator, VertexLabelAnnotator, ) -from supervision.keypoint.core import KeyPoints +from supervision.key_points.core import KeyPoints from supervision.metrics.detection import ConfusionMatrix, MeanAveragePrecision from supervision.tracker.byte_tracker.core import ByteTrack from supervision.utils.conversion import cv2_to_pillow, pillow_to_cv2 diff --git a/supervision/keypoint/__init__.py b/supervision/key_points/__init__.py similarity index 100% rename from supervision/keypoint/__init__.py rename to supervision/key_points/__init__.py diff --git a/supervision/keypoint/annotators.py b/supervision/key_points/annotators.py similarity index 99% rename from supervision/keypoint/annotators.py rename to supervision/key_points/annotators.py index d83d7d5e..ab9f04d1 100644 --- a/supervision/keypoint/annotators.py +++ b/supervision/key_points/annotators.py @@ -11,8 +11,8 @@ from supervision.detection.utils.boxes import pad_boxes, spread_out_boxes from supervision.draw.color import Color from supervision.draw.utils import draw_rounded_rectangle from supervision.geometry.core import Rect -from supervision.keypoint.core import KeyPoints -from supervision.keypoint.skeletons import SKELETONS_BY_VERTEX_COUNT +from supervision.key_points.core import KeyPoints +from supervision.key_points.skeletons import SKELETONS_BY_VERTEX_COUNT from supervision.utils.conversion import ensure_cv2_image_for_annotation diff --git a/supervision/keypoint/core.py b/supervision/key_points/core.py similarity index 81% rename from supervision/keypoint/core.py rename to supervision/key_points/core.py index 0d57c561..fb37e628 100644 --- a/supervision/keypoint/core.py +++ b/supervision/key_points/core.py @@ -10,7 +10,7 @@ import numpy.typing as npt from supervision.config import CLASS_NAME_DATA_FIELD from supervision.detection.core import Detections from supervision.detection.utils.internal import get_data_item, is_data_equal -from supervision.validators import validate_keypoints_fields +from supervision.validators import validate_key_points_fields @dataclass @@ -23,7 +23,7 @@ class KeyPoints: === "Ultralytics" - Use [`sv.KeyPoints.from_ultralytics`](/latest/keypoint/core/#supervision.keypoint.core.KeyPoints.from_ultralytics) + Use [`sv.KeyPoints.from_ultralytics`](/latest/keypoint/core/#supervision.key_points.core.KeyPoints.from_ultralytics) method, which accepts [YOLOv8-pose](https://docs.ultralytics.com/models/yolov8/), [YOLO11-pose](https://docs.ultralytics.com/models/yolo11/) [pose](https://docs.ultralytics.com/tasks/pose/) result. @@ -41,7 +41,7 @@ class KeyPoints: === "Inference" - Use [`sv.KeyPoints.from_inference`](/latest/keypoint/core/#supervision.keypoint.core.KeyPoints.from_inference) + Use [`sv.KeyPoints.from_inference`](/latest/keypoint/core/#supervision.key_points.core.KeyPoints.from_inference) method, which accepts [Inference](https://inference.roboflow.com/) pose result. ```python @@ -58,7 +58,7 @@ class KeyPoints: === "MediaPipe" - Use [`sv.KeyPoints.from_mediapipe`](/latest/keypoint/core/#supervision.keypoint.core.KeyPoints.from_mediapipe) + Use [`sv.KeyPoints.from_mediapipe`](/latest/keypoint/core/#supervision.key_points.core.KeyPoints.from_mediapipe) method, which accepts [MediaPipe](https://github.com/google-ai-edge/mediapipe) pose result. @@ -88,11 +88,62 @@ class KeyPoints: key_points = sv.KeyPoints.from_mediapipe( pose_landmarker_result, (image_width, image_height)) ``` + + === "Transformers" + + Use [`sv.KeyPoints.from_transformers`](/latest/keypoint/core/#supervision.key_points.core.KeyPoints.from_transformers) + method, which accepts [ViTPose](https://huggingface.co/docs/transformers/en/model_doc/vitpose) result. + + ```python + from PIL import Image + import requests + import supervision as sv + import torch + from transformers import ( + AutoProcessor, + RTDetrForObjectDetection, + VitPoseForPoseEstimation, + ) + + device = "cuda" if torch.cuda.is_available() else "cpu" + image = Image.open() + + DETECTION_MODEL_ID = "PekingU/rtdetr_r50vd_coco_o365" + + detection_processor = AutoProcessor.from_pretrained(DETECTION_MODEL_ID, use_fast=True) + detection_model = RTDetrForObjectDetection.from_pretrained(DETECTION_MODEL_ID, device_map=DEVICE) + + inputs = detection_processor(images=frame, return_tensors="pt").to(DEVICE) + + with torch.no_grad(): + outputs = detection_model(**inputs) + + target_size = torch.tensor([(frame.height, frame.width)]) + results = detection_processor.post_process_object_detection( + outputs, target_sizes=target_size, threshold=0.3) + + detections = sv.Detections.from_transformers(results[0]) + boxes = sv.xyxy_to_xywh(detections[detections.class_id == 0].xyxy) + + POSE_ESTIMATION_MODEL_ID = "usyd-community/vitpose-base-simple" + + pose_estimation_processor = AutoProcessor.from_pretrained(POSE_ESTIMATION_MODEL_ID) + pose_estimation_model = VitPoseForPoseEstimation.from_pretrained( + POSE_ESTIMATION_MODEL_ID, device_map=DEVICE) + + inputs = pose_estimation_processor(frame, boxes=[boxes], return_tensors="pt").to(DEVICE) + + with torch.no_grad(): + outputs = pose_estimation_model(**inputs) + + results = pose_estimation_processor.post_process_pose_estimation(outputs, boxes=[boxes]) + key_point = sv.KeyPoints.from_transformers(results[0]) + ``` Attributes: xy (np.ndarray): An array of shape `(n, m, 2)` containing `n` detected objects, each composed of `m` equally-sized - sets of keypoints, where each point is `[x, y]`. + sets of key points, where each point is `[x, y]`. class_id (Optional[np.ndarray]): An array of shape `(n,)` containing the class ids of the detected objects. confidence (Optional[np.ndarray]): An array of shape @@ -109,7 +160,7 @@ class KeyPoints: data: dict[str, npt.NDArray[Any] | list] = field(default_factory=dict) def __post_init__(self): - validate_keypoints_fields( + validate_key_points_fields( xy=self.xy, confidence=self.confidence, class_id=self.class_id, @@ -514,13 +565,13 @@ class KeyPoints: return cls.empty() @classmethod - def from_transformers(cls, transfomers_results: Any) -> KeyPoints: + def from_transformers(cls, transformers_results: Any) -> KeyPoints: """ Create a `sv.KeyPoints` object from the [Transformers](https://github.com/huggingface/transformers) inference result. Args: - transfomers_results (Any): The output of a + transformers_results (Any): The output of a Transformers model containing instances with prediction data. Returns: @@ -576,8 +627,8 @@ class KeyPoints: """ # noqa: E501 // docs - if "keypoints" in transfomers_results[0]: - if transfomers_results[0]["keypoints"].cpu().numpy().size == 0: + if "keypoints" in transformers_results[0]: + if transformers_results[0]["keypoints"].cpu().numpy().size == 0: return cls.empty() result_data = [ @@ -585,7 +636,7 @@ class KeyPoints: result["keypoints"].cpu().numpy(), result["scores"].cpu().numpy(), ) - for result in transfomers_results + for result in transformers_results ] xy, scores = zip(*result_data) @@ -599,55 +650,72 @@ class KeyPoints: return cls.empty() def __getitem__( - self, index: int | slice | list[int] | np.ndarray | str - ) -> KeyPoints | list | np.ndarray | None: - """ - Get a subset of the `sv.KeyPoints` object or access an item from its data field. - - When provided with an integer, slice, list of integers, or a numpy array, this - method returns a new `sv.KeyPoints` object that represents a subset of the - original `sv.KeyPoints`. When provided with a string, it accesses the - corresponding item in the data dictionary. - - Args: - index (Union[int, slice, List[int], np.ndarray, str]): The index, indices, - or key to access a subset of the `sv.KeyPoints` or an item from the - data. - - Returns: - A subset of the `sv.KeyPoints` object or an item from the data field. - - Examples: - ```python - import supervision as sv - - key_points = sv.KeyPoints() - - # access the first keypoint using an integer index - key_points[0] - - # access the first 10 keypoints using index slice - key_points[0:10] - - # access selected keypoints using a list of indices - key_points[[0, 2, 4]] - - # access keypoints with selected class_id - key_points[key_points.class_id == 0] - - # access keypoints with confidence greater than 0.5 - key_points[key_points.confidence > 0.5] - ``` - """ + self, index: int | slice | list[int] | np.ndarray | tuple | str + ) -> KeyPoints | np.ndarray | list | None: if isinstance(index, str): return self.data.get(index) - if isinstance(index, int): - index = [index] + + if not isinstance(index, tuple): + index = (index, slice(None)) + + i, j = index + + if isinstance(i, int): + i = [i] + + if isinstance(i, list) and all(isinstance(x, bool) for x in i): + i = np.array(i) + if isinstance(j, list) and all(isinstance(x, bool) for x in j): + j = np.array(j) + + if isinstance(i, np.ndarray) and i.dtype == bool: + i = np.flatnonzero(i) + if isinstance(j, np.ndarray) and j.dtype == bool: + j = np.flatnonzero(j) + + if ( + isinstance(i, (list, np.ndarray)) + and isinstance(j, (list, np.ndarray)) + and not np.isscalar(i) + and not np.isscalar(j) + ): + i, j = np.ix_(i, j) + + xy_selected = self.xy[i, j] + + conf_selected = ( + self.confidence[i, j] if self.confidence is not None else None + ) + + class_id_selected = ( + self.class_id[i] if self.class_id is not None else None + ) + + data_selected = get_data_item(self.data, i) + + if xy_selected.ndim == 1: + xy_selected = xy_selected.reshape(1, 1, 2) + if conf_selected is not None: + conf_selected = conf_selected.reshape(1, 1) + elif xy_selected.ndim == 2: + if np.isscalar(index[0]) or ( + isinstance(index[0], np.ndarray) and index[0].ndim == 0 + ): + xy_selected = xy_selected[np.newaxis, ...] + if conf_selected is not None: + conf_selected = conf_selected[np.newaxis, ...] + elif np.isscalar(index[1]) or ( + isinstance(index[1], np.ndarray) and index[1].ndim == 0 + ): + xy_selected = xy_selected[:, np.newaxis, :] + if conf_selected is not None: + conf_selected = conf_selected[:, np.newaxis] + return KeyPoints( - xy=self.xy[index], - confidence=self.confidence[index] if self.confidence is not None else None, - class_id=self.class_id[index] if self.class_id is not None else None, - data=get_data_item(self.data, index), + xy=xy_selected, + confidence=conf_selected, + class_id=class_id_selected, + data=data_selected, ) def __setitem__(self, key: str, value: np.ndarray | list): @@ -668,12 +736,12 @@ class KeyPoints: model = YOLO('yolov8s.pt') result = model(image)[0] - keypoints = sv.KeyPoints.from_ultralytics(result) + key_points = sv.KeyPoints.from_ultralytics(result) - keypoints['class_name'] = [ + key_points['class_name'] = [ model.model.names[class_id] for class_id - in keypoints.class_id + in key_points.class_id ] ``` """ @@ -688,7 +756,7 @@ class KeyPoints: @classmethod def empty(cls) -> KeyPoints: """ - Create an empty Keypoints object with no keypoints. + Create an empty KeyPoints object with no key points. Returns: An empty `sv.KeyPoints` object. @@ -706,9 +774,9 @@ class KeyPoints: """ Returns `True` if the `KeyPoints` object is considered empty. """ - empty_keypoints = KeyPoints.empty() - empty_keypoints.data = self.data - return self == empty_keypoints + empty_key_points = KeyPoints.empty() + empty_key_points.data = self.data + return self == empty_key_points def as_detections( self, selected_keypoint_indices: Iterable[int] | None = None @@ -716,21 +784,21 @@ class KeyPoints: """ Convert a KeyPoints object to a Detections object. This approximates the bounding box of the detected object by - taking the bounding box that fits all keypoints. + taking the bounding box that fits all key points. Arguments: selected_keypoint_indices (Optional[Iterable[int]]): The - indices of the keypoints to include in the bounding box - calculation. This helps focus on a subset of keypoints, - e.g. when some are occluded. Captures all keypoints by default. + indices of the key points to include in the bounding box + calculation. This helps focus on a subset of key points, + e.g. when some are occluded. Captures all key points by default. Returns: detections (Detections): The converted detections object. Examples: ```python - keypoints = sv.KeyPoints.from_inference(...) - detections = keypoints.as_detections() + key_points = sv.KeyPoints.from_inference(...) + detections = key_points.as_detections() ``` """ if self.is_empty(): diff --git a/supervision/keypoint/skeletons.py b/supervision/key_points/skeletons.py similarity index 100% rename from supervision/keypoint/skeletons.py rename to supervision/key_points/skeletons.py diff --git a/supervision/validators/__init__.py b/supervision/validators/__init__.py index 97fedabd..f051d89d 100644 --- a/supervision/validators/__init__.py +++ b/supervision/validators/__init__.py @@ -53,7 +53,7 @@ def validate_confidence(confidence: Any, n: int) -> None: ) -def validate_keypoint_confidence(confidence: Any, n: int, m: int) -> None: +def validate_key_point_confidence(confidence: Any, n: int, m: int) -> None: expected_shape = f"({n, m})" actual_shape = str(getattr(confidence, "shape", None)) @@ -126,7 +126,7 @@ def validate_detections_fields( validate_data(data, n) -def validate_keypoints_fields( +def validate_key_points_fields( xy: Any, class_id: Any, confidence: Any, @@ -136,7 +136,7 @@ def validate_keypoints_fields( m = len(xy[0]) if len(xy) > 0 else 0 validate_xy(xy, n, m) validate_class_id(class_id, n) - validate_keypoint_confidence(confidence, n, m) + validate_key_point_confidence(confidence, n, m) validate_data(data, n) diff --git a/test/key_points/__init__.py b/test/key_points/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/test/key_points/test_core.py b/test/key_points/test_core.py new file mode 100644 index 00000000..f12cc2fd --- /dev/null +++ b/test/key_points/test_core.py @@ -0,0 +1,444 @@ +import numpy as np +import pytest +from contextlib import nullcontext as DoesNotRaise +from supervision.key_points.core import KeyPoints +from test.test_utils import mock_key_points + +KEY_POINTS = mock_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), + mock_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), + mock_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, + mock_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, + mock_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], + mock_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]), + mock_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], + mock_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]), + mock_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)), + mock_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)), + mock_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), + mock_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), + mock_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]), + mock_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])), + mock_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]), + mock_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])), + mock_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), + mock_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), + mock_key_points( + xy=[ + [[8, 9]], + ], + confidence=[ + [0.5], + ], + class_id=[0], + ), + DoesNotRaise(), + ), # select the last anchor from the first skeleton by index + ], +) +def test_key_points_getitem(key_points, index, expected_result, exception): + with exception: + result = key_points[index] + assert result == expected_result diff --git a/test/test_utils.py b/test/test_utils.py index 19fffad5..0a97bf4b 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -3,14 +3,13 @@ from __future__ import annotations from typing import Any import numpy as np -import numpy.typing as npt from supervision.detection.core import Detections -from supervision.keypoint.core import KeyPoints +from supervision.key_points.core import KeyPoints def mock_detections( - xyxy: npt.NDArray[np.float32], + xyxy: list[list[float]], mask: list[np.ndarray] | None = None, confidence: list[float] | None = None, class_id: list[int] | None = None, @@ -34,9 +33,9 @@ def mock_detections( ) -def mock_keypoints( - xy: npt.NDArray[np.float32], - confidence: list[float] | None = None, +def mock_key_points( + xy: list[list[list[float]]], + confidence: list[list[float]] | None = None, class_id: list[int] | None = None, data: dict[str, list[Any]] | None = None, ) -> KeyPoints: