From 422ca868638cecbfd785b3260128cb1e0a38dfd3 Mon Sep 17 00:00:00 2001 From: Miguelmelon Date: Fri, 14 Jun 2024 18:38:44 +0200 Subject: [PATCH] Chi-squared instead of Euclidean for HoC --- people_tracking_v2/scripts/comparison_node.py | 12 ++++++++---- people_tracking_v2/scripts/decision_node.py | 2 +- people_tracking_v2/scripts/yolo_seg.py | 8 ++++---- people_tracking_v2/tools/recorder.py | 17 +++++++++++++---- 4 files changed, 26 insertions(+), 13 deletions(-) diff --git a/people_tracking_v2/scripts/comparison_node.py b/people_tracking_v2/scripts/comparison_node.py index 592b6f6..281db30 100755 --- a/people_tracking_v2/scripts/comparison_node.py +++ b/people_tracking_v2/scripts/comparison_node.py @@ -126,13 +126,17 @@ def normalize_vector(self, vector): return norm_vector / norm_sum def compute_hoc_distance_score(self, hue_vector, sat_vector, val_vector): - """Compute the distance score between the current detection and saved data (HoC).""" - hue_distance = self.compute_distance(hue_vector, self.operator_hue_avg) - sat_distance = self.compute_distance(sat_vector, self.operator_sat_avg) - val_distance = self.compute_distance(val_vector, self.operator_val_avg) + """Compute the Chi-Squared distance score between the current detection and saved data (HoC).""" + hue_distance = self.compute_chi_squared_distance(hue_vector, self.operator_hue_avg) + sat_distance = self.compute_chi_squared_distance(sat_vector, self.operator_sat_avg) + val_distance = self.compute_chi_squared_distance(val_vector, self.operator_val_avg) return (hue_distance + sat_distance + val_distance) + def compute_chi_squared_distance(self, vector1, vector2): + """Compute the Chi-Squared distance between two vectors.""" + return 0.5 * np.sum(((vector1 - vector2) ** 2) / (vector1 + vector2 + 1e-10)) # Adding a small value to avoid division by zero + def compute_distance(self, vector1, vector2): """Compute the Euclidean distance between two vectors (General).""" return np.linalg.norm(vector1 - vector2) diff --git a/people_tracking_v2/scripts/decision_node.py b/people_tracking_v2/scripts/decision_node.py index 9d6b3ae..4b392b3 100755 --- a/people_tracking_v2/scripts/decision_node.py +++ b/people_tracking_v2/scripts/decision_node.py @@ -27,7 +27,7 @@ def sync_callback(self, comparison_msg, detection_msg): """Callback function to handle synchronized comparison scores and detection info.""" # Define thresholds iou_threshold = 0.8 - hoc_threshold = 0.05 + hoc_threshold = 0.3 pose_threshold = 0.5 # Example threshold for pose distance iou_detections = [] diff --git a/people_tracking_v2/scripts/yolo_seg.py b/people_tracking_v2/scripts/yolo_seg.py index 21700a6..8ad1127 100755 --- a/people_tracking_v2/scripts/yolo_seg.py +++ b/people_tracking_v2/scripts/yolo_seg.py @@ -18,10 +18,10 @@ sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'src', 'people_tracking')) from kalman_filter import KalmanFilterCV # Import the Kalman Filter class -laptop = True # sys.argv[1] -name_subscriber_RGB = 'Webcam/image_raw' # if laptop == "True" else '/hero/head_rgbd_sensor/rgb/image_raw' -depth_camera = False # if sys.argv[2] == "False" else True -save_data = False # if sys.argv[3] == "False" else True +laptop = sys.argv[1] +name_subscriber_RGB = 'Webcam/image_raw' if laptop == "True" else '/hero/head_rgbd_sensor/rgb/image_raw' +depth_camera = False if sys.argv[2] == "False" else True +save_data = False if sys.argv[3] == "False" else True class YoloSegNode: def __init__(self): diff --git a/people_tracking_v2/tools/recorder.py b/people_tracking_v2/tools/recorder.py index 20e656c..b18c258 100644 --- a/people_tracking_v2/tools/recorder.py +++ b/people_tracking_v2/tools/recorder.py @@ -4,9 +4,10 @@ from sensor_msgs.msg import Image from cv_bridge import CvBridge, CvBridgeError import cv2 +import os class VideoRecorder: - def __init__(self): + def __init__(self, save_path): rospy.init_node('video_recorder', anonymous=True) # Subscribe to the RGB and depth image topics @@ -22,6 +23,11 @@ def __init__(self): self.fourcc = cv2.VideoWriter_fourcc(*'XVID') self.fps = 30 # You can adjust this based on your requirements + # Ensure the save path exists + self.save_path = save_path + if not os.path.exists(self.save_path): + os.makedirs(self.save_path) + def rgb_image_callback(self, data): try: # Convert the ROS Image message to an OpenCV image @@ -30,7 +36,8 @@ def rgb_image_callback(self, data): # Initialize the RGB video writer once we get the first image if self.rgb_out is None: self.rgb_frame_size = (rgb_image.shape[1], rgb_image.shape[0]) - self.rgb_out = cv2.VideoWriter('rgb_output.avi', self.fourcc, self.fps, self.rgb_frame_size) + rgb_filename = os.path.join(self.save_path, 'rgb_output.avi') + self.rgb_out = cv2.VideoWriter(rgb_filename, self.fourcc, self.fps, self.rgb_frame_size) # Write the RGB image to the video file self.rgb_out.write(rgb_image) @@ -53,7 +60,8 @@ def depth_image_callback(self, data): # Initialize the depth video writer once we get the first image if self.depth_out is None: self.depth_frame_size = (depth_image_colored.shape[1], depth_image_colored.shape[0]) - self.depth_out = cv2.VideoWriter('depth_output.avi', self.fourcc, self.fps, self.depth_frame_size) + depth_filename = os.path.join(self.save_path, 'depth_output.avi') + self.depth_out = cv2.VideoWriter(depth_filename, self.fourcc, self.fps, self.depth_frame_size) # Write the depth image to the video file self.depth_out.write(depth_image_colored) @@ -73,7 +81,8 @@ def cleanup(self): cv2.destroyAllWindows() if __name__ == '__main__': - recorder = VideoRecorder() + save_path = '~/hero_videos' # Replace with your desired save path + recorder = VideoRecorder(save_path) try: rospy.spin() except KeyboardInterrupt: