diff --git a/docs/how_to/track_objects.md b/docs/how_to/track_objects.md index 464f6b8d..784cb7bf 100644 --- a/docs/how_to/track_objects.md +++ b/docs/how_to/track_objects.md @@ -6,7 +6,7 @@ comments: true Leverage Supervision's advanced capabilities for enhancing your video analysis by seamlessly [tracking](/latest/trackers/) objects recognized by -a multitude of object detection and segmentation models. This comprehensive guide will +a multitude of object detection, segmentation and keypoint models. This comprehensive guide will take you through the steps to perform inference using the YOLOv8 model via either the [Inference](https://github.com/roboflow/inference) or [Ultralytics](https://github.com/ultralytics/ultralytics) packages. Following this, @@ -21,6 +21,7 @@ example. You can do this using from supervision.assets import download_assets, VideoAssets download_assets(VideoAssets.PEOPLE_WALKING) +download_assets(VideoAssets.SKIING) ``` +## Tracking Key Points + +Keypoint tracking is currently supported via the conversion of `KeyPoints` to `Detections`. This is achieved with the [`keypoints_to_detections`](/latest/utils/datatypes/#supervision.utils.datatypes.keypoints_to_detections) function. We'll use a different video as well as [`DetectionsSmoother`](/latest/detection/tools/smoother/) to stabilize the boxes. + +!!! tip + + You may use the `selected_keypoint_indices` argument to specify a subset of keypoints to convert. This is useful when some keypoints could be occluded. For example: a person might swing their arm, causing the elbow to be occluded by the torso sometimes. + +=== "Ultralytics" + + ```{ .py hl_lines="5 7 14-15 17 33" } + import numpy as np + import supervision as sv + from ultralytics import YOLO + + model = YOLO("yolov8m-pose.pt") + tracker = sv.ByteTrack() + smoother = sv.DetectionsSmoother() + box_annotator = sv.BoundingBoxAnnotator() + label_annotator = sv.LabelAnnotator() + trace_annotator = sv.TraceAnnotator() + + def callback(frame: np.ndarray, _: int) -> np.ndarray: + results = model(frame)[0] + keypoints = sv.KeyPoints.from_ultralytics(results) + detections = sv.keypoints_to_detections(keypoints) + detections = tracker.update_with_detections(detections) + detections = smoother.update_with_detections(detections) + + labels = [ + f"#{tracker_id} {results.names[class_id]}" + for class_id, tracker_id + in zip(detections.class_id, detections.tracker_id) + ] + + annotated_frame = box_annotator.annotate( + frame.copy(), detections=detections) + annotated_frame = label_annotator.annotate( + annotated_frame, detections=detections, labels=labels) + return trace_annotator.annotate( + annotated_frame, detections=detections) + + sv.process_video( + source_path="skiing.mp4", + target_path="result.mp4", + callback=callback + ) + ``` + +=== "Inference" + + ```{ .py hl_lines="5-6 8 15-16 18 34" } + import numpy as np + import supervision as sv + from inference.models.utils import get_roboflow_model + + model = get_roboflow_model( + model_id="yolov8m-pose-640", api_key=) + tracker = sv.ByteTrack() + smoother = sv.DetectionsSmoother() + box_annotator = sv.BoundingBoxAnnotator() + label_annotator = sv.LabelAnnotator() + trace_annotator = sv.TraceAnnotator() + + def callback(frame: np.ndarray, _: int) -> np.ndarray: + results = model.infer(frame)[0] + keypoints = sv.KeyPoints.from_inference(results) + detections = sv.keypoints_to_detections(keypoints) + detections = tracker.update_with_detections(detections) + detections = smoother.update_with_detections(detections) + + labels = [ + f"#{tracker_id} {results.names[class_id]}" + for class_id, tracker_id + in zip(detections.class_id, detections.tracker_id) + ] + + annotated_frame = box_annotator.annotate( + frame.copy(), detections=detections) + annotated_frame = label_annotator.annotate( + annotated_frame, detections=detections, labels=labels) + return trace_annotator.annotate( + annotated_frame, detections=detections) + + sv.process_video( + source_path="skiing.mp4", + target_path="result.mp4", + callback=callback + ) + ``` + + + This structured walkthrough should give a detailed pathway to annotate videos effectively using Supervision’s various functionalities, including object tracking and trace annotations. 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/detection/tools/polygon_zone.py b/supervision/detection/tools/polygon_zone.py index f69f3c9f..5cd976b1 100644 --- a/supervision/detection/tools/polygon_zone.py +++ b/supervision/detection/tools/polygon_zone.py @@ -19,7 +19,7 @@ class PolygonZone: !!! warning - LineZone uses the `tracker_id`. Read + PolygonZone uses the `tracker_id`. Read [here](/latest/trackers/) to learn how to plug tracking into your inference pipeline. diff --git a/supervision/keypoint/core.py b/supervision/keypoint/core.py index 36d6a596..252fb63f 100644 --- a/supervision/keypoint/core.py +++ b/supervision/keypoint/core.py @@ -612,3 +612,11 @@ class KeyPoints: ``` """ return cls(xy=np.empty((0, 0, 2), dtype=np.float32)) + + def is_empty(self) -> bool: + """ + Returns `True` if the `KeyPoints` object is considered empty. + """ + empty_keypoints = KeyPoints.empty() + empty_keypoints.data = self.data + return self == empty_keypoints diff --git a/supervision/utils/datatypes.py b/supervision/utils/datatypes.py new file mode 100644 index 00000000..6f2e8879 --- /dev/null +++ b/supervision/utils/datatypes.py @@ -0,0 +1,72 @@ +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] + + # [0, 0] used by some frameworks to indicate missing keypoints + xy = xy[~np.all(xy == 0, axis=1)] + if len(xy) == 0: + xyxy = np.array([[0, 0, 0, 0]], dtype=np.float32) + else: + 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 + detections = detections[detections.area > 0] + + return detections