From 644af4b977ac69f7868ad5b1f66129535e273557 Mon Sep 17 00:00:00 2001 From: Miguelmelon Date: Sun, 26 May 2024 15:40:34 +0200 Subject: [PATCH] Included operator HoC data recorder --- people_tracking_v2/CMakeLists.txt | 1 + people_tracking_v2/scripts/HoC.py | 31 ++----- people_tracking_v2/scripts/trial.py | 89 +++++++++++++++++++ ...launcher.launch => ppl_tracking_v2.launch} | 0 .../launchers/save-HoC-data.launch | 34 +++++++ .../{scripts => tools}/plotter.py | 0 people_tracking_v2/tools/save-HoC-data.py | 51 +++++++++++ 7 files changed, 182 insertions(+), 24 deletions(-) create mode 100644 people_tracking_v2/scripts/trial.py rename people_tracking_v2/src/people_tracking/launchers/{yolo_launcher.launch => ppl_tracking_v2.launch} (100%) create mode 100644 people_tracking_v2/src/people_tracking/launchers/save-HoC-data.launch rename people_tracking_v2/{scripts => tools}/plotter.py (100%) create mode 100755 people_tracking_v2/tools/save-HoC-data.py diff --git a/people_tracking_v2/CMakeLists.txt b/people_tracking_v2/CMakeLists.txt index 384d37f..3ffe43c 100644 --- a/people_tracking_v2/CMakeLists.txt +++ b/people_tracking_v2/CMakeLists.txt @@ -41,6 +41,7 @@ install(PROGRAMS scripts/yolo_seg.py scripts/HoC.py src/kalman_filter.py + tools/save-HoC-data.py DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION} ) diff --git a/people_tracking_v2/scripts/HoC.py b/people_tracking_v2/scripts/HoC.py index a16dcd0..dfa70bf 100755 --- a/people_tracking_v2/scripts/HoC.py +++ b/people_tracking_v2/scripts/HoC.py @@ -5,24 +5,17 @@ from cv_bridge import CvBridge, CvBridgeError import cv2 import numpy as np -import os class HoCNode: - def __init__(self): - rospy.init_node('hoc_node', anonymous=True) + def __init__(self, initialize_node=True): + if initialize_node: + rospy.init_node('hoc_node', anonymous=True) self.bridge = CvBridge() self.segmented_images_sub = rospy.Subscriber('/segmented_images', SegmentedImages, self.segmented_images_callback) - # Ensure the directory exists - self.hoc_data_dir = os.path.expanduser('~/hoc_data') - if not os.path.exists(self.hoc_data_dir): - os.makedirs(self.hoc_data_dir) - rospy.loginfo(f"Created directory: {self.hoc_data_dir}") - else: - rospy.loginfo(f"Using existing directory: {self.hoc_data_dir}") - - rospy.spin() + if initialize_node: + rospy.spin() def segmented_images_callback(self, msg): rospy.loginfo(f"Received batch of {len(msg.images)} segmented images") @@ -30,18 +23,8 @@ def segmented_images_callback(self, msg): try: segmented_image = self.bridge.imgmsg_to_cv2(segmented_image_msg, "bgr8") hoc_hue, hoc_sat = self.compute_hoc(segmented_image) - #rospy.loginfo(f'HoC for segmented image #{i + 1} (Hue): {hoc_hue}') - #rospy.loginfo(f'HoC for segmented image #{i + 1} (Saturation): {hoc_sat}') - - # Save the HoC data - hue_save_path = os.path.join(self.hoc_data_dir, f'hoc_hue_detection_{i + 1}.npy') - sat_save_path = os.path.join(self.hoc_data_dir, f'hoc_sat_detection_{i + 1}.npy') - np.save(hue_save_path, hoc_hue) - np.save(sat_save_path, hoc_sat) - - # Print statements to verify file saving - #rospy.loginfo(f'Saved Hue HoC to {hue_save_path}') - #rospy.loginfo(f'Saved Sat HoC to {sat_save_path}') + rospy.loginfo(f'Computed HoC for segmented image #{i + 1}') + # You can process hoc_hue and hoc_sat here or pass them to another node except CvBridgeError as e: rospy.logerr(f"Failed to convert segmented image: {e}") diff --git a/people_tracking_v2/scripts/trial.py b/people_tracking_v2/scripts/trial.py new file mode 100644 index 0000000..e298d71 --- /dev/null +++ b/people_tracking_v2/scripts/trial.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python +import rospy +import csv +import numpy as np +from people_tracking_v2.msg import DetectionArray +from std_msgs.msg import String +from collections import namedtuple + +# Named tuple for storing target information +Target = namedtuple("Target", ["nr_batch", "time", "colour_vector"]) + +class TargetMemoryNode: + def __init__(self): + # Initialize the ROS node + rospy.init_node('target_memory', anonymous=True) + + # Subscriber to HoC detections + self.subscriber_hoc = rospy.Subscriber('/hero/HoC', DetectionArray, self.hoc_callback) + + # Publisher for debug information or status updates + self.publisher_debug = rospy.Publisher('/hero/target_memory/debug', String, queue_size=10) + + # Dictionary to store targets with their batch number as key + self.targets = {} + + # Load saved HoC data + self.load_hoc_data() + + rospy.spin() + + def load_hoc_data(self): + """Load the saved HoC data from the CSV file.""" + try: + with open('hoc_data.csv', 'r') as csvfile: + csv_reader = csv.reader(csvfile) + for row in csv_reader: + nr_batch = int(row[0]) + time = rospy.Time(float(row[1])) + colour_vector = list(map(float, row[2:])) + self.targets[nr_batch] = Target(nr_batch=nr_batch, time=time, colour_vector=colour_vector) + rospy.loginfo("Loaded HoC data from hoc_data.csv") + except Exception as e: + rospy.logerr(f"Failed to load HoC data: {e}") + + def hoc_callback(self, msg): + """Callback function to handle new HoC detections.""" + for detection in msg.detections: + similarity_score = self.compute_similarity(detection) + self.update_target(detection, similarity_score) + + def update_target(self, detection, similarity_score): + """Update or add a target based on the detection.""" + batch_number = detection.nr_batch + colour_vector = detection.colour_vector + + if batch_number in self.targets: + # Update existing target + existing_target = self.targets[batch_number] + updated_target = Target(nr_batch=batch_number, time=existing_target.time, colour_vector=colour_vector) + self.targets[batch_number] = updated_target + rospy.loginfo(f"Updated target batch {batch_number} with similarity score {similarity_score:.2f}") + else: + # Add new target + new_target = Target(nr_batch=batch_number, time=rospy.Time.now(), colour_vector=colour_vector) + self.targets[batch_number] = new_target + rospy.loginfo(f"Added new target batch {batch_number} with similarity score {similarity_score:.2f}") + + # Publish debug information + debug_msg = String() + debug_msg.data = f"Total targets: {len(self.targets)} | Last similarity score: {similarity_score:.2f}" + self.publisher_debug.publish(debug_msg) + + def compute_similarity(self, detection): + """Compute the similarity score between the current detection and saved targets.""" + max_similarity = 0 + current_vector = np.array(detection.colour_vector) + + for target in self.targets.values(): + saved_vector = np.array(target.colour_vector) + similarity = np.dot(current_vector, saved_vector) / (np.linalg.norm(current_vector) * np.linalg.norm(saved_vector)) + max_similarity = max(max_similarity, similarity) + + return max_similarity + +if __name__ == '__main__': + try: + TargetMemoryNode() + except rospy.ROSInterruptException: + pass diff --git a/people_tracking_v2/src/people_tracking/launchers/yolo_launcher.launch b/people_tracking_v2/src/people_tracking/launchers/ppl_tracking_v2.launch similarity index 100% rename from people_tracking_v2/src/people_tracking/launchers/yolo_launcher.launch rename to people_tracking_v2/src/people_tracking/launchers/ppl_tracking_v2.launch diff --git a/people_tracking_v2/src/people_tracking/launchers/save-HoC-data.launch b/people_tracking_v2/src/people_tracking/launchers/save-HoC-data.launch new file mode 100644 index 0000000..45716e1 --- /dev/null +++ b/people_tracking_v2/src/people_tracking/launchers/save-HoC-data.launch @@ -0,0 +1,34 @@ + + + + + + + + + + + + + + \ No newline at end of file diff --git a/people_tracking_v2/scripts/plotter.py b/people_tracking_v2/tools/plotter.py similarity index 100% rename from people_tracking_v2/scripts/plotter.py rename to people_tracking_v2/tools/plotter.py diff --git a/people_tracking_v2/tools/save-HoC-data.py b/people_tracking_v2/tools/save-HoC-data.py new file mode 100755 index 0000000..ea298e8 --- /dev/null +++ b/people_tracking_v2/tools/save-HoC-data.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python + +import rospy +import numpy as np +from people_tracking_v2.msg import SegmentedImages # Custom message for batch segmented images +from cv_bridge import CvBridge, CvBridgeError +import os +import sys + +# Add the scripts directory to the Python path +sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'scripts')) + +from HoC import HoCNode # Import the HoCNode class + +class SaveHoCDataNode: + def __init__(self): + rospy.init_node('save_hoc_data_node', anonymous=True) + + self.bridge = CvBridge() + self.segmented_images_sub = rospy.Subscriber('/segmented_images', SegmentedImages, self.segmented_images_callback) + + # File to save HoC data + self.hoc_data_file = os.path.expanduser('~/hoc_data/hoc_data.npz') + + # Instantiate the HoCNode for HoC computation without initializing the node again + self.hoc_node = HoCNode(initialize_node=False) + + rospy.spin() + + def segmented_images_callback(self, msg): + rospy.loginfo(f"Received batch of {len(msg.images)} segmented images") + all_hue_histograms = [] + all_sat_histograms = [] + for i, segmented_image_msg in enumerate(msg.images): + try: + segmented_image = self.bridge.imgmsg_to_cv2(segmented_image_msg, "bgr8") + hoc_hue, hoc_sat = self.hoc_node.compute_hoc(segmented_image) + all_hue_histograms.append(hoc_hue) + all_sat_histograms.append(hoc_sat) + except CvBridgeError as e: + rospy.logerr(f"Failed to convert segmented image: {e}") + + # Save all histograms in a single .npz file + np.savez(self.hoc_data_file, hue=all_hue_histograms, sat=all_sat_histograms) + rospy.loginfo(f'Saved HoC data to {self.hoc_data_file}') + +if __name__ == '__main__': + try: + SaveHoCDataNode() + except rospy.ROSInterruptException: + pass