diff --git a/docs/utils/datatypes.md b/docs/utils/datatypes.md new file mode 100644 index 00000000..5e0560bd --- /dev/null +++ b/docs/utils/datatypes.md @@ -0,0 +1,12 @@ +--- +comments: true +status: new +--- + +# Data Types Utils + +
+

keypoints_to_detections

+
+ +:::supervision.utils.datatypes.keypoints_to_detections diff --git a/mkdocs.yml b/mkdocs.yml index b30dbcfc..8c939cf0 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -79,6 +79,7 @@ nav: - File: utils/file.md - Draw: utils/draw.md - Geometry: utils/geometry.md + - Datatypes: utils/datatypes.md - Assets: assets.md - Cookbooks: cookbooks.md - Cheatsheet: https://roboflow.github.io/cheatsheet-supervision/ diff --git a/supervision/__init__.py b/supervision/__init__.py index 746b2f67..04813ef2 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -100,6 +100,7 @@ from supervision.keypoint.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 +from supervision.utils.datatypes import keypoints_to_detections from supervision.utils.file import list_files_with_extensions from supervision.utils.image import ( ImageSink, diff --git a/supervision/utils/datatypes.py b/supervision/utils/datatypes.py new file mode 100644 index 00000000..dafbc4cb --- /dev/null +++ b/supervision/utils/datatypes.py @@ -0,0 +1,65 @@ +from typing import Iterable, Optional + +import numpy as np + +from supervision.detection.core import Detections +from supervision.keypoint.core import KeyPoints + + +def keypoints_to_detections( + keypoints: KeyPoints, selected_keypoint_indices: Optional[Iterable[int]] = None +) -> Detections: + """ + 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. + + Arguments: + keypoints (KeyPoints): The keypoints to convert to detections. + 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. + + Returns: + detections (Detections): The converted detections object. + + Example: + ```python + keypoints = sv.KeyPoints.from_inference(...) + detections = keypoints_to_detections(keypoints) + ``` + """ + if keypoints.is_empty(): + return Detections.empty() + + detections_list = [] + for i, xy in enumerate(keypoints.xy): + if selected_keypoint_indices: + xy = xy[selected_keypoint_indices] + x_min = xy[:, 0].min() + x_max = xy[:, 0].max() + y_min = xy[:, 1].min() + y_max = xy[:, 1].max() + xyxy = np.array([[x_min, y_min, x_max, y_max]], dtype=np.float32) + + if keypoints.confidence is None: + confidence = None + else: + confidence = keypoints.confidence[i] + if selected_keypoint_indices: + confidence = confidence[selected_keypoint_indices] + confidence = np.array([confidence.mean()], dtype=np.float32) + + detections_list.append( + Detections( + xyxy=xyxy, + confidence=confidence, + ) + ) + + detections = Detections.merge(detections_list) + detections.class_id = keypoints.class_id + detections.data = keypoints.data + + return detections