Merge pull request #1658 from roboflow/feat/keypoints-to-detections-and-keypoint-tracking

`keypoints_to_detections` and Keypoint tracking

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)
 ```
 
 <video controls>
@@ -41,6 +42,7 @@ This `callback` function will be essential in the subsequent steps of the tutori
 it will be modified to include tracking, labeling, and trace annotations.
 
 === "Ultralytics"
+
     ```{ .py }
     import numpy as np
     import supervision as sv
@@ -62,6 +64,7 @@ it will be modified to include tracking, labeling, and trace annotations.
     ```
 
 === "Inference"
+
     ```{ .py }
     import numpy as np
     import supervision as sv
@@ -95,6 +98,7 @@ functionality, each detected object is assigned a unique tracker ID,
 enabling the continuous following of the object's motion path across different frames.
 
 === "Ultralytics"
+
     ```{ .py hl_lines="6 12" }
     import numpy as np
     import supervision as sv
@@ -118,6 +122,7 @@ enabling the continuous following of the object's motion path across different f
     ```
 
 === "Inference"
+
     ```{ .py hl_lines="6 12" }
     import numpy as np
     import supervision as sv
@@ -149,6 +154,7 @@ in Supervision, we can overlay the tracker IDs and class labels on the detected
 offering a clear visual representation of each object's class and unique identifier.
 
 === "Ultralytics"
+
     ```{ .py hl_lines="8 15-19 23-24" }
     import numpy as np
     import supervision as sv
@@ -183,6 +189,7 @@ offering a clear visual representation of each object's class and unique identif
     ```
 
 === "Inference"
+
     ```{ .py hl_lines="8 15-19 23-24" }
     import numpy as np
     import supervision as sv
@@ -229,6 +236,7 @@ allows for visualizing the trajectories of objects, helping in understanding the
 movement patterns and interactions between objects in the video.
 
 === "Ultralytics"
+
     ```{ .py hl_lines="9 26-27" }
     import numpy as np
     import supervision as sv
@@ -266,6 +274,7 @@ movement patterns and interactions between objects in the video.
     ```
 
 === "Inference"
+
     ```{ .py hl_lines="9 26-27" }
     import numpy as np
     import supervision as sv
@@ -306,6 +315,101 @@ movement patterns and interactions between objects in the video.
     <source src="https://media.roboflow.com/supervision/video-examples/how-to/track-objects/annotate-video-with-traces.mp4" type="video/mp4">
 </video>
 
+## 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=<ROBOFLOW 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
+    )
+    ```
+
+<video controls>
+    <source src="https://media.roboflow.com/supervision/video-examples/how-to/track-objects/track-keypoints.mp4" type="video/mp4">
+</video>
+
 This structured walkthrough should give a detailed pathway to annotate videos
 effectively using Supervision’s various functionalities, including object tracking and
 trace annotations.

docs/utils/datatypes.md

@@ -0,0 +1,12 @@
+---
+comments: true
+status: new
+---
+
+# Data Types Utils
+
+<div class="md-typeset">
+    <h2><a href="#supervision.utils.datatypes.keypoints_to_detections">keypoints_to_detections</a></h2>
+</div>
+
+:::supervision.utils.datatypes.keypoints_to_detections

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/

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,

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.
 

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

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