(pose) migrate body parts to separate file

image_recognition_pose_estimation/src/image_recognition_pose_estimation/__init__.py

@@ -0,0 +1 @@
+from . import body_parts

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",
+    "left-eye": "left-eye",
+    "right-eye": "right-eye",
+    "left-ear": "left-ear",
+    "right-ear": "right-ear",
+    "left-shoulder": "left-shoulder",
+    "right-shoulder": "right-shoulder",
+    "left-elbow": "left-elbow",
+    "right-elbow": "right-elbow",
+    "left-wrist": "left-wrist",
+    "right-wrist": "right-wrist",
+    "left-hip": "left-hip",
+    "right-hip": "right-hip",
+    "left-knee": "left-knee",
+    "right-knee": "right-knee",
+    "left-ankle": "left-ankle",
+    "right-ankle": "right-ankle",
+}
+
+BODY_PART_LINKS = [
+    # The lowest index first
+    # Matches the keys of BODY_PARTS
+    # HEAD
+    ("nose", "left-eye"),
+    ("left-eye", "left-ear"),
+    ("nose", "right-eye"),
+    ("left-eye", "right-eye"),
+    ("right-eye", "right-ear"),
+    # Left side
+    ("left-ear", "left-shoulder"),
+    ("left-shoulder", "left-elbow"),
+    ("left-elbow", "left-wrist"),
+    ("left-shoulder", "left-hip"),
+    ("left-hip", "left-knee"),
+    ("left-knee", "left-ankle"),
+
+    # Right side
+    ("right-ear", "right-shoulder"),
+    ("right-shoulder", "right-elbow"),
+    ("right-elbow", "right-wrist"),
+    ("right-shoulder", "right-hip"),
+    ("right-hip", "right-knee"),
+    ("right-knee", "right-ankle"),
+]

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$")
 
@@ -111,7 +73,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),
                                     )
                                 ]