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..1ce75b5e --- /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", + "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"), +] 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..8598530f 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$") @@ -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), ) ]