From cab06f2c2b424e547b7c4e80f142b0ba61aeb70d Mon Sep 17 00:00:00 2001
From: Linas Kondrackis <linas.ko+dev@skiff.com>
Date: Thu, 25 Apr 2024 12:29:08 +0300
Subject: [PATCH] Implementation of from_yolo_nas for keypoints

* Missing batch image support
---
 supervision/keypoint/core.py      | 55 ++++++++++++++++++++++++++++---
 supervision/keypoint/skeletons.py | 24 ++++++++++++++
 2 files changed, 74 insertions(+), 5 deletions(-)

diff --git a/supervision/keypoint/core.py b/supervision/keypoint/core.py
index 8a97b51c..587d459c 100644
--- a/supervision/keypoint/core.py
+++ b/supervision/keypoint/core.py
@@ -50,7 +50,8 @@ class KeyPoints:
     xy: npt.NDArray[np.float32]
     class_id: Optional[npt.NDArray[np.int_]] = None
     confidence: Optional[npt.NDArray[np.float32]] = None
-    data: Dict[str, Union[npt.NDArray[Any], List]] = field(default_factory=dict)
+    data: Dict[str, Union[npt.NDArray[Any], List]
+               ] = field(default_factory=dict)
 
     def __post_init__(self):
         validate_keypoints_fields(
@@ -103,7 +104,7 @@ class KeyPoints:
     @classmethod
     def from_ultralytics(cls, ultralytics_results) -> KeyPoints:
         """
-        Creates a Keypoints instance from a
+        Creates a KeyPoints instance from a
             [YOLOv8](https://github.com/ultralytics/ultralytics) inference result.
 
         Args:
@@ -111,7 +112,7 @@ class KeyPoints:
                 The output Results instance from YOLOv8
 
         Returns:
-            KeyPoints: A new Keypoints object.
+            KeyPoints: A new KeyPoints object.
 
         Example:
             ```python
@@ -130,12 +131,56 @@ class KeyPoints:
 
         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])
+        class_names = np.array([ultralytics_results.names[i]
+                               for i in class_id])
 
         confidence = ultralytics_results.keypoints.conf.cpu().numpy()
         data = {CLASS_NAME_DATA_FIELD: class_names}
         return cls(xy, class_id, confidence, data)
 
+    @classmethod
+    def from_yolo_nas(cls, yolo_nas_results) -> KeyPoints:
+        """
+        Create a KeyPoints instance from a YOLO NAS results.
+
+        Args:
+            yolo_nas_results (ImagePoseEstimationPrediction): The output object from YOLO NAS.
+
+        Returns:
+            KeyPoints: A new KeyPoints object.
+
+        Example:
+            ```python
+            import cv2
+            import supervision as sv
+            import super_gradients
+
+            image = cv2.imread(<SOURCE_IMAGE_PATH>)
+
+            yolo_nas = super_gradients.training.models.get(
+                "yolo_nas_pose_s", pretrained_weights="coco_pose").to("cuda")
+
+            results = yolo_nas.predict(image, conf=0.1)
+            keypoints = sv.KeyPoints.from_yolo_nas(results)
+            ```
+        """
+        if len(yolo_nas_results.prediction.poses) == 0:
+            return cls.empty()
+
+        # TODO: multi-image input
+
+        xy = yolo_nas_results.prediction.poses[:, :, :2]
+        confidence = yolo_nas_results.prediction.poses[:, :, 2]
+        class_id = [0] * len(xy)
+        data = {CLASS_NAME_DATA_FIELD: [f"person" for _ in xy]}
+
+        return cls(
+            xy=xy,
+            confidence=confidence,
+            class_id=np.array(class_id),
+            data=data,
+        )
+
     def __getitem__(
         self, index: Union[int, slice, List[int], np.ndarray, str]
     ) -> Union["KeyPoints", List, np.ndarray, None]:
@@ -216,7 +261,7 @@ class KeyPoints:
 
         self.data[key] = value
 
-    @classmethod
+    @ classmethod
     def empty(cls) -> KeyPoints:
         """
         Create an empty Keypoints object with no keypoints.
diff --git a/supervision/keypoint/skeletons.py b/supervision/keypoint/skeletons.py
index 6c110854..01f768b1 100644
--- a/supervision/keypoint/skeletons.py
+++ b/supervision/keypoint/skeletons.py
@@ -25,6 +25,30 @@ class Skeleton(Enum):
         (17, 15),
     ]
 
+    # Hardcoding here, but it also comes within model results
+    # Note: the classes in the model are 0-indexed, while the skeleton is 1-indexed
+    YOLO_NAS = [
+        (0, 1),
+        (0, 2),
+        (1, 2),
+        (1, 3),
+        (2, 4),
+        (3, 5),
+        (4, 6),
+        (5, 6),
+        (5, 7),
+        (5, 11),
+        (6, 8),
+        (6, 12),
+        (7, 9),
+        (8, 10),
+        (11, 12),
+        (11, 13),
+        (12, 14),
+        (13, 15),
+        (14, 16)
+    ]
+
 
 SKELETONS_BY_EDGE_COUNT: Dict[int, Edges] = {}
 SKELETONS_BY_VERTEX_COUNT: Dict[int, Edges] = {}