Operator pose data recorder

people_tracking_v2/CMakeLists.txt

@@ -19,6 +19,7 @@ add_message_files(
   DetectionArray.msg
   SegmentedImages.msg
   HoCVector.msg
+  BodySize.msg
 )
 
 generate_messages(
@@ -42,7 +43,8 @@ install(PROGRAMS
   scripts/HoC.py
   scripts/comparison_node.py
   src/kalman_filter.py
-  tools/save-HoC-data.py
+  tools/save_HoC_data.py
+  tools/save_pose_data.py
   DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION}
 )
 

people_tracking_v2/msg/BodySize.msg

@@ -0,0 +1,3 @@
+Header header
+float32 left_shoulder_hip_distance
+float32 right_shoulder_hip_distance

people_tracking_v2/scripts/pose_estimation_node.py

@@ -1,4 +1,4 @@
-#! /usr/bin/env python
+#!/usr/bin/env python
 
 import os
 import socket
@@ -18,6 +18,7 @@
 from sensor_msgs.msg import Image
 
 from people_tracking.yolo_pose_wrapper import YoloPoseWrapper
+from people_tracking_v2.msg import PoseDistance
 
 class PoseEstimationNode:
     def __init__(
@@ -48,7 +49,7 @@ def __init__(
         # uses a thread per pub/sub and service. Since the openpose wrapper is created in the main thread, we have
         # to communicate our openpose requests (inputs) to the main thread where the request is processed by the
         # openpose wrapper (step 1).
-        # We have a separate spin loop in the main thread that checks whether there are items in the input q and
+        # We have a separate spin loop in the main thead that checks whether there are items in the input q and
         # processes these using the Openpose wrapper (step 2).
         # When the processing has finished, we add the result in the corresponding output queue (specified by the
         # request in the input queue) (step 3).
@@ -68,6 +69,7 @@ def __init__(
         self._recognize_srv = rospy.Service("recognize", Recognize, self._recognize_srv)
         self._image_subscriber = rospy.Subscriber("/Webcam/image_raw", Image, self._image_callback)
         self._recognitions_publisher = rospy.Publisher("/pose_recognitions", Recognitions, queue_size=10)
+        self._pose_distance_publisher = rospy.Publisher("/pose_distances", PoseDistance, queue_size=10)
         if self._topic_publish_result_image or self._service_publish_result_image:
             self._result_image_publisher = rospy.Publisher("/pose_result_image", Image, queue_size=10)
 
@@ -87,16 +89,20 @@ def _image_callback(self, image_msg):
         self._input_q.put((image_msg, self._topic_save_images, self._topic_publish_result_image, False))
         recognitions, result_image, pose_details = self._wrapper.detect_poses(self._bridge.imgmsg_to_cv2(image_msg, "bgr8"))
 
-        # Calculate distances
+        # Calculate distances and publish them
         for pose in pose_details:
             try:
+                pose_distance_msg = PoseDistance()
+                pose_distance_msg.header.stamp = rospy.Time.now()
                 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}")
+                    pose_distance_msg.left_shoulder_hip_distance = self._wrapper.compute_distance(pose["LShoulder"], pose["LHip"])
+                    rospy.loginfo(f"Left Shoulder-Hip Distance: {pose_distance_msg.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}")
+                    pose_distance_msg.right_shoulder_hip_distance = self._wrapper.compute_distance(pose["RShoulder"], pose["RHip"])
+                    rospy.loginfo(f"Right Shoulder-Hip Distance: {pose_distance_msg.right_shoulder_hip_distance:.2f}")
+
+                self._pose_distance_publisher.publish(pose_distance_msg)
             except Exception as e:
                 rospy.logerr(f"Error computing distance: {e}")
 
@@ -142,13 +148,17 @@ def _get_recognitions(self, image_msg, save_images, publish_images):
         # Calculate distances and log them
         for pose in pose_details:
             try:
+                pose_distance_msg = PoseDistance()
+                pose_distance_msg.header.stamp = rospy.Time.now()
                 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}")
+                    pose_distance_msg.left_shoulder_hip_distance = self._wrapper.compute_distance(pose["LShoulder"], pose["LHip"])
+                    rospy.loginfo(f"Left Shoulder-Hip Distance: {pose_distance_msg.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}")
+                    pose_distance_msg.right_shoulder_hip_distance = self._wrapper.compute_distance(pose["RShoulder"], pose["RHip"])
+                    rospy.loginfo(f"Right Shoulder-Hip Distance: {pose_distance_msg.right_shoulder_hip_distance:.2f}")
+
+                self._pose_distance_publisher.publish(pose_distance_msg)
             except Exception as e:
                 rospy.logerr(f"Error computing distance: {e}")
 

people_tracking_v2/tools/save_pose_data.py

@@ -0,0 +1,60 @@
+#!/usr/bin/env python
+
+import rospy
+import numpy as np
+from people_tracking_v2.msg import PoseDistance
+import os
+
+class SavePoseDataNode:
+    def __init__(self):
+        # Initialize the ROS node
+        rospy.init_node('save_pose_data_node', anonymous=True)
+
+        # Subscriber to pose distance data
+        self.pose_sub = rospy.Subscriber('/pose_distances', PoseDistance, self.pose_callback)
+
+        # File to save pose data
+        self.pose_data_file = os.path.expanduser('~/pose_data/pose_data.npz')
+
+        # Ensure the directory exists
+        pose_data_dir = os.path.dirname(self.pose_data_file)
+        if not os.path.exists(pose_data_dir):
+            os.makedirs(pose_data_dir)
+            rospy.loginfo(f"Created directory: {pose_data_dir}")
+        else:
+            rospy.loginfo(f"Using existing directory: {pose_data_dir}")
+
+        self.pose_data = {
+            'left_shoulder_hip_distance': [],
+            'right_shoulder_hip_distance': []
+        }
+
+        rospy.spin()
+
+    def pose_callback(self, msg):
+        rospy.loginfo("Received pose distance data")
+
+        # Append the distances to the data dictionary
+        self.pose_data['left_shoulder_hip_distance'].append(msg.left_shoulder_hip_distance)
+        self.pose_data['right_shoulder_hip_distance'].append(msg.right_shoulder_hip_distance)
+        rospy.loginfo(f"Left Shoulder-Hip Distance: {msg.left_shoulder_hip_distance:.2f}")
+        rospy.loginfo(f"Right Shoulder-Hip Distance: {msg.right_shoulder_hip_distance:.2f}")
+
+        # Save the pose data periodically or based on some condition
+        if len(self.pose_data['left_shoulder_hip_distance']) >= 10:  # Save every 10 messages as an example
+            self.save_pose_data()
+
+    def save_pose_data(self):
+        """Save the collected pose data to a .npz file."""
+        np.savez(self.pose_data_file, **self.pose_data)
+        rospy.loginfo(f"Saved pose data to {self.pose_data_file}")
+        self.pose_data = {
+            'left_shoulder_hip_distance': [],
+            'right_shoulder_hip_distance': []
+        }  # Clear the data after saving
+
+if __name__ == '__main__':
+    try:
+        SavePoseDataNode()
+    except rospy.ROSInterruptException:
+        pass