diff --git a/supervision/keypoints/annotate.py b/supervision/keypoints/annotate.py new file mode 100644 index 00000000..4c37b248 --- /dev/null +++ b/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 diff --git a/supervision/keypoints/core.py b/supervision/keypoints/core.py index 646b99db..25f14fae 100644 --- a/supervision/keypoints/core.py +++ b/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), + ] +)