From 1a8e7f0ced647ed48656a6d1d6e27d484c2962d9 Mon Sep 17 00:00:00 2001
From: Matthijs van der Burgh <MatthijsBurgh@outlook.com>
Date: Tue, 12 Mar 2024 21:13:09 +0100
Subject: [PATCH] (pose) migrate body parts to separate file

---
 .../__init__.py                               |  1 +
 .../body_parts.py                             | 64 +++++++++++++++++++
 .../yolo_pose_wrapper.py                      | 44 ++-----------
 3 files changed, 69 insertions(+), 40 deletions(-)
 create mode 100644 image_recognition_pose_estimation/src/image_recognition_pose_estimation/body_parts.py

diff --git a/image_recognition_pose_estimation/src/image_recognition_pose_estimation/__init__.py b/image_recognition_pose_estimation/src/image_recognition_pose_estimation/__init__.py
index e69de29b..24d70ed7 100644
--- a/image_recognition_pose_estimation/src/image_recognition_pose_estimation/__init__.py
+++ b/image_recognition_pose_estimation/src/image_recognition_pose_estimation/__init__.py
@@ -0,0 +1 @@
+from . import body_parts
diff --git a/image_recognition_pose_estimation/src/image_recognition_pose_estimation/body_parts.py b/image_recognition_pose_estimation/src/image_recognition_pose_estimation/body_parts.py
new file mode 100644
index 00000000..c392ba05
--- /dev/null
+++ b/image_recognition_pose_estimation/src/image_recognition_pose_estimation/body_parts.py
@@ -0,0 +1,64 @@
+# Yolo pose keypoint labels
+# 0: nose
+# 1: left-eye
+# 2: right-eye
+# 3: left-ear
+# 4: right-ear
+# 5: left-shoulder
+# 6: right-shoulder
+# 7: left-elbow
+# 8: right-elbow
+# 9: left-wrist
+# 10: right-wrist
+# 11: left-hip
+# 12: right-hip
+# 13: left-knee
+# 14: right-knee
+# 15: left-ankle
+# 16: right-ankle
+
+BODY_PARTS = {
+    "Nose": "Nose",
+    "LEye": "LEye",
+    "REye": "REye",
+    "LEar": "LEar",
+    "REar": "REar",
+    "LShoulder": "LShoulder",
+    "RShoulder": "RShoulder",
+    "LElbow": "LElbow",
+    "RElbow": "RElbow",
+    "LWrist": "LWrist",
+    "RWrist": "RWrist",
+    "LHip": "LHip",
+    "RHip": "RHip",
+    "LKnee": "LKnee",
+    "RKnee": "RKnee",
+    "LAnkle": "LAnkle",
+    "RAnkle": "RAnkle",
+}
+
+BODY_PART_LINKS = [
+    # The lowest index first
+    # Matches the keys of BODY_PARTS
+    # HEAD
+    ("Nose", "LEye"),
+    ("LEye", "LEar"),
+    ("Nose", "REye"),
+    ("LEye", "REye"),
+    ("REye", "REar"),
+    # Left side
+    ("LEar", "LShoulder"),
+    ("LShoulder", "LElbow"),
+    ("LElbow", "LWrist"),
+    ("LShoulder", "LHip"),
+    ("LHip", "LKnee"),
+    ("LKnee", "LAnkle"),
+
+    # Right side
+    ("REar", "RShoulder"),
+    ("RShoulder", "RElbow"),
+    ("RElbow", "RWrist"),
+    ("RShoulder", "RHip"),
+    ("RHip", "RKnee"),
+    ("RKnee", "RAnkle"),
+]
diff --git a/image_recognition_pose_estimation/src/image_recognition_pose_estimation/yolo_pose_wrapper.py b/image_recognition_pose_estimation/src/image_recognition_pose_estimation/yolo_pose_wrapper.py
index d188cd46..fd7ca025 100644
--- a/image_recognition_pose_estimation/src/image_recognition_pose_estimation/yolo_pose_wrapper.py
+++ b/image_recognition_pose_estimation/src/image_recognition_pose_estimation/yolo_pose_wrapper.py
@@ -5,49 +5,11 @@
 import numpy as np
 import torch
 from image_recognition_msgs.msg import CategoricalDistribution, CategoryProbability, Recognition
+from image_recognition_pose_estimation.body_parts import BODY_PARTS
 from sensor_msgs.msg import RegionOfInterest
 from ultralytics import YOLO
 from ultralytics.engine.results import Results
 
-# Yolo pose keypoint labels
-# 0: nose
-# 1: left-eye
-# 2: right-eye
-# 3: left-ear
-# 4: right-ear
-# 5: left-shoulder
-# 6: right-shoulder
-# 7: left-elbow
-# 8: right-elbow
-# 9: left-wrist
-# 10: right-wrist
-# 11: left-hip
-# 12: right-hip
-# 13: left-knee
-# 14: right-knee
-# 15: left-ankle
-# 16: right-ankle
-
-YOLO_POSE_KEYPOINT_LABELS = [
-    "nose",
-    "left-eye",
-    "right-eye",
-    "left-ear",
-    "right-ear",
-    "left-shoulder",
-    "right-shoulder",
-    "left-elbow",
-    "right-elbow",
-    "left-wrist",
-    "right-wrist",
-    "left-hip",
-    "right-hip",
-    "left-knee",
-    "right-knee",
-    "left-ankle",
-    "right-ankle",
-]
-
 
 YOLO_POSE_PATTERN = re.compile(r"^yolov8(?:([nsml])|(x))-pose(?(2)-p6|)?.pt$")
 
@@ -101,6 +63,8 @@ def detect_poses(self, image: np.ndarray, conf: float = 0.25) -> Tuple[List[Reco
         result = results[0]  # Only using
         overlayed_image = result.plot(boxes=False)
 
+        body_parts = list(BODY_PARTS.values())
+
         for i, person in enumerate(result.keypoints.cpu().numpy()):
             for j, (x, y, pred_conf) in enumerate(person.data[0]):
                 if pred_conf > 0 and x > 0 and y > 0:
@@ -111,7 +75,7 @@ def detect_poses(self, image: np.ndarray, conf: float = 0.25) -> Tuple[List[Reco
                             categorical_distribution=CategoricalDistribution(
                                 probabilities=[
                                     CategoryProbability(
-                                        label=YOLO_POSE_KEYPOINT_LABELS[j],
+                                        label=body_parts[j],
                                         probability=float(pred_conf),
                                     )
                                 ]