diff --git a/people_tracking_v2/scripts/pose_estimation_node.py b/people_tracking_v2/scripts/pose_estimation_node.py index 6a4c475..b889e0d 100755 --- a/people_tracking_v2/scripts/pose_estimation_node.py +++ b/people_tracking_v2/scripts/pose_estimation_node.py @@ -89,13 +89,16 @@ def _image_callback(self, image_msg): # Calculate distances for pose in pose_details: - if "LShoulder" in pose and "RShoulder" in pose: - shoulder_distance = self._wrapper.compute_distance(pose["LShoulder"], pose["RShoulder"]) - rospy.loginfo(f"Shoulder Distance: {shoulder_distance:.2f}") + try: + if "LShoulder" in pose and "LHip" in pose: + left_shoulder_hip_distance = self._wrapper.compute_distance(pose["LShoulder"], pose["LHip"]) + rospy.loginfo(f"Left Shoulder-Hip Distance: {left_shoulder_hip_distance:.2f}") - if "LHip" in pose and "RHip" in pose: - hip_distance = self._wrapper.compute_distance(pose["LHip"], pose["RHip"]) - rospy.loginfo(f"Hip Distance: {hip_distance:.2f}") + if "RShoulder" in pose and "RHip" in pose: + right_shoulder_hip_distance = self._wrapper.compute_distance(pose["RShoulder"], pose["RHip"]) + rospy.loginfo(f"Right Shoulder-Hip Distance: {right_shoulder_hip_distance:.2f}") + except Exception as e: + rospy.logerr(f"Error computing distance: {e}") self._recognitions_publisher.publish( Recognitions(header=image_msg.header, recognitions=self._subscriber_output_q.get()) @@ -136,15 +139,18 @@ def _get_recognitions(self, image_msg, save_images, publish_images): if publish_images: self._result_image_publisher.publish(self._bridge.cv2_to_imgmsg(result_image, "bgr8")) - # Calculate distances + # Calculate distances and log them for pose in pose_details: - if "LShoulder" in pose and "LHip" in pose: - left_shoulder_hip_distance = self._wrapper.compute_distance(pose["LShoulder"], pose["LHip"]) - rospy.loginfo(f"Left Shoulder-Hip Distance: {left_shoulder_hip_distance:.2f}") - - if "RShoulder" in pose and "RHip" in pose: - right_shoulder_hip_distance = self._wrapper.compute_distance(pose["RShoulder"], pose["RHip"]) - rospy.loginfo(f"Right Shoulder-Hip Distance: {right_shoulder_hip_distance:.2f}") + try: + if "LShoulder" in pose and "LHip" in pose: + left_shoulder_hip_distance = self._wrapper.compute_distance(pose["LShoulder"], pose["LHip"]) + rospy.loginfo(f"Left Shoulder-Hip Distance: {left_shoulder_hip_distance:.2f}") + + if "RShoulder" in pose and "RHip" in pose: + right_shoulder_hip_distance = self._wrapper.compute_distance(pose["RShoulder"], pose["RHip"]) + rospy.loginfo(f"Right Shoulder-Hip Distance: {right_shoulder_hip_distance:.2f}") + except Exception as e: + rospy.logerr(f"Error computing distance: {e}") return recognitions diff --git a/people_tracking_v2/src/people_tracking/yolo_pose_wrapper.py b/people_tracking_v2/src/people_tracking/yolo_pose_wrapper.py index e471352..61f7698 100644 --- a/people_tracking_v2/src/people_tracking/yolo_pose_wrapper.py +++ b/people_tracking_v2/src/people_tracking/yolo_pose_wrapper.py @@ -10,7 +10,6 @@ from ultralytics import YOLO from ultralytics.engine.results import Results - YOLO_POSE_PATTERN = re.compile(r"^yolov8(?:([nsml])|(x))-pose(?(2)-p6|)?.pt$") ALLOWED_DEVICE_TYPES = ["cpu", "cuda"] @@ -52,23 +51,23 @@ def __init__(self, model_name: str = "yolov8n-pose.pt", device: str = "cuda:0", def detect_poses(self, image: np.ndarray, conf: float = 0.25) -> Tuple[List[Recognition], np.ndarray, List[dict]]: # Detect poses - # This is a wrapper of predict, but we might want to use track results: List[Results] = self._model(image, conf=conf) # Accepts a list if not results: return [], image, [] recognitions = [] - pose_details = [] - result = results[0] # Only using + result = results[0] # Only using the first result overlayed_image = result.plot(boxes=False) body_parts = list(BODY_PARTS.values()) + pose_details = [] # Extracted pose details for calculating distances for i, person in enumerate(result.keypoints.cpu().numpy()): pose = {} for j, (x, y, pred_conf) in enumerate(person.data[0]): if pred_conf > 0 and x > 0 and y > 0: + pose[body_parts[j]] = (x, y) recognitions.append( Recognition( group_id=i, @@ -83,12 +82,10 @@ def detect_poses(self, image: np.ndarray, conf: float = 0.25) -> Tuple[List[Reco ), ) ) - pose[body_parts[j]] = (x, y) pose_details.append(pose) return recognitions, overlayed_image, pose_details - @staticmethod def compute_distance(self, point1, point2): """Compute the Euclidean distance between two points.""" return np.sqrt((point1[0] - point2[0]) ** 2 + (point1[1] - point2[1]) ** 2)