Add keypoints skeleton, 2x annotators

supervision/keypoints/annotate.py

@@ -0,0 +1,102 @@
+from abc import ABC, abstractmethod
+
+import cv2
+import numpy as np
+
+from supervision.annotators.base import ImageType
+from supervision.annotators.utils import scene_to_annotator_img_type
+from supervision.draw.color import Color
+from supervision.keypoints.core import KeyPoints, Skeleton
+
+
+class BaseKeyPointAnnotator(ABC):
+    @abstractmethod
+    def annotate(self, scene: ImageType, detections: KeyPoints) -> ImageType:
+        pass
+
+
+class PointAnnotator(BaseKeyPointAnnotator):
+    def __init__(
+        self,
+        color: Color = Color.GREEN,
+        radius: int = 4,
+    ) -> None:
+        """
+        Most basic keypoint annotator.
+        """
+        self.color = color
+        self.radius = radius
+
+    @scene_to_annotator_img_type
+    def annotate(self, scene: ImageType, keypoints: KeyPoints) -> ImageType:
+        # TODO: I'm getting shape [1, N, 2] here - not [N, 2].
+        #       Seems like it accounts for more than 1 person in an image.
+        if len(keypoints) == 0:
+            return scene
+
+        xy = keypoints.xy[0]
+        for i, (x, y) in enumerate(xy):
+            cv2.circle(
+                img=scene,
+                center=(int(x), int(y)),
+                radius=self.radius,
+                color=self.color.as_bgr(),
+                thickness=-1,
+            )
+
+        return scene
+
+
+class SkeletonAnnotator(BaseKeyPointAnnotator):
+    def __init__(
+        self,
+        point_color: Color = Color.GREEN,
+        point_radius: int = 4,
+        limb_color: Color = Color.GREEN,
+        limb_thickness: int = 2,
+    ) -> None:
+        self.point_color = point_color
+        self.point_radius = point_radius
+        self.limb_color = limb_color
+        self.limb_thickness = limb_thickness
+
+    @scene_to_annotator_img_type
+    def annotate(
+        self, scene: ImageType, keypoints: KeyPoints, skeleton: Skeleton
+    ) -> ImageType:
+        if len(keypoints) == 0:
+            return scene
+        if keypoints.class_id is None:
+            raise ValueError("KeyPoints must have class_id to annotate a skeleton")
+
+        # print(keypoints)
+
+        xy = keypoints.xy[0]
+        class_id = keypoints.class_id[0]
+
+        for i, (x, y) in enumerate(xy):
+            cv2.circle(
+                img=scene,
+                center=(int(x), int(y)),
+                radius=self.point_radius,
+                color=self.point_color.as_bgr(),
+                thickness=-1,
+            )
+
+        for class_a, class_b in skeleton.limbs:
+            xy_a = xy[class_a - 1]
+            xy_b = xy[class_b - 1]
+            missing_a = np.allclose(xy_a, 0)
+            missing_b = np.allclose(xy_b, 0)
+            if missing_a or missing_b:
+                continue
+
+            cv2.line(
+                img=scene,
+                pt1=(int(xy_a[0]), int(xy_a[1])),
+                pt2=(int(xy_b[0]), int(xy_b[1])),
+                color=self.limb_color.as_bgr(),
+                thickness=self.limb_thickness,
+            )
+
+        return scene

supervision/keypoints/core.py

@@ -132,14 +132,15 @@ class KeyPoints:
             result = model(image)[0]
             keypoints = sv.KeyPoints.from_ultralytics(result)
             ```
-        """  #
+        """
+        if len(ultralytics_results.keypoints.xy) == 0:
+            return cls.empty()
 
+        xy = ultralytics_results.keypoints.xy.cpu().numpy()
         class_id = ultralytics_results.boxes.cls.cpu().numpy().astype(int)
         class_names = np.array([ultralytics_results.names[i] for i in class_id])
 
-        xy = ultralytics_results.keypoints.xy.cpu().numpy()
         confidence = ultralytics_results.keypoints.conf.cpu().numpy()
-        class_id = ultralytics_results.boxes.cls.cpu().numpy().astype(int)
         data = {CLASS_NAME_DATA_FIELD: class_names}
         return cls(xy, class_id, confidence, data)
 
@@ -238,3 +239,39 @@ class KeyPoints:
             ```
         """
         return cls(xy=np.empty((0, 0, 2), dtype=np.float32))
+
+
+class Skeleton:
+    """
+    The `Skeleton` class connects keypoints to form a skeleton.
+
+    It provides utility methods to compute angles within itself, compare to skeletons, etc.
+    """
+
+    limbs: List[Tuple[int, int]]
+
+    def __init__(self, limbs: List[Tuple[int, int]]):
+        self.limbs = limbs
+
+
+YOLO_V8_SKELETON = Skeleton(
+    [
+        (1, 2),
+        (1, 3),
+        (2, 3),
+        (2, 4),
+        (3, 5),
+        (6, 12),
+        (6, 7),
+        (6, 8),
+        (7, 13),
+        (7, 9),
+        (8, 10),
+        (9, 11),
+        (12, 13),
+        (14, 12),
+        (15, 13),
+        (16, 14),
+        (17, 15),
+    ]
+)