From 3ee8ffa4dbc902a21b26bff495cb7d4261ea2f39 Mon Sep 17 00:00:00 2001 From: Miguelmelon <m.colera.borras@student.tue.nl> Date: Sun, 16 Jun 2024 18:38:14 +0200 Subject: [PATCH] Fixing KF 5 --- .../scripts/pose_estimation_node.py | 5 + people_tracking_v2/scripts/yolo_seg.py | 102 ++++++++++-------- 2 files changed, 62 insertions(+), 45 deletions(-) diff --git a/people_tracking_v2/scripts/pose_estimation_node.py b/people_tracking_v2/scripts/pose_estimation_node.py index e58b811..b4b4b83 100755 --- a/people_tracking_v2/scripts/pose_estimation_node.py +++ b/people_tracking_v2/scripts/pose_estimation_node.py @@ -39,6 +39,11 @@ def image_callback(self, image_msg): for detection in self.current_detections: if detection.iou > self.iou_threshold: + # Set the height of the specific detection to -1 if IoU threshold is passed + pose_distance_msg = BodySize() + pose_distance_msg.id = detection.id + pose_distance_msg.head_feet_distance = -1 # Set to -1 for detections with high IoU + pose_distance_array.distances.append(pose_distance_msg) rospy.loginfo(f"Skipping detection ID {detection.id} due to high IoU value with operator") continue diff --git a/people_tracking_v2/scripts/yolo_seg.py b/people_tracking_v2/scripts/yolo_seg.py index f4a1b74..e53c453 100755 --- a/people_tracking_v2/scripts/yolo_seg.py +++ b/people_tracking_v2/scripts/yolo_seg.py @@ -95,14 +95,18 @@ def image_callback(self, data): # Import Depth Image if depth_camera: - cv_depth_image = self.bridge.imgmsg_to_cv2(self.depth_images[-1], desired_encoding='passthrough') + try: + cv_depth_image = self.bridge.imgmsg_to_cv2(self.depth_images[-1], desired_encoding='passthrough') + except (CvBridgeError, IndexError) as e: + rospy.logerr(f"CV Bridge Error or Index Error: {e}") + cv_depth_image = None else: cv_depth_image = None # Save RGB and Depth Images if required if save_data: cv2.imwrite(f"{self.save_path}{self.batch_nr}.png", cv_image) - if depth_camera: + if depth_camera and cv_depth_image is not None: cv_depth_image_path = f"{self.save_path}{self.batch_nr}_depth.png" cv2.imwrite(cv_depth_image_path, cv_depth_image) @@ -183,6 +187,12 @@ def image_callback(self, data): # Publish individual segmented images self.individual_segmented_image_pub.publish(segmented_image_msg) + # Draw a red dot at the center of the detection if it's the operator + if detection.id == self.operator_id: + x_center = int((detection.x1 + detection.x2) / 2) + y_center = int((detection.y1 + detection.y2) / 2) + cv2.circle(bounding_box_image, (x_center, y_center), 5, (0, 0, 255), -1) + # Initialize the operator using the detection with the specified operator ID if not self.operator_initialized and self.operator_id is not None: for detection in detection_array.detections: @@ -200,46 +210,51 @@ def image_callback(self, data): self.kalman_filter_operator.predict() x_pred, y_pred = self.kalman_filter_operator.get_state()[:2] - # Ensure operator_box is not None before using it - if operator_box is not None: - # Draw predicted position and bounding box - box_width = operator_box[2] - operator_box[0] - box_height = operator_box[3] - operator_box[1] - x_pred1, y_pred1 = int(x_pred - box_width / 2), int(y_pred - box_height / 2) - x_pred2, y_pred2 = int(x_pred + box_width / 2), int(y_pred + box_height / 2) + # Use default box dimensions if operator_box is None + box_width = 50 if operator_box is None else operator_box[2] - operator_box[0] + box_height = 100 if operator_box is None else operator_box[3] - operator_box[1] - cv2.rectangle(bounding_box_image, (x_pred1, y_pred1), (x_pred2, y_pred2), (255, 0, 0), 2) # Blue box - cv2.circle(bounding_box_image, (int(x_pred), int(y_pred)), 5, (255, 0, 0), -1) + x_pred1, y_pred1 = int(x_pred - box_width / 2), int(y_pred - box_height / 2) + x_pred2, y_pred2 = int(x_pred + box_width / 2), int(y_pred + box_height / 2) - # Calculate IoU for all detections with the operator's predicted bounding box - best_iou = -1 - best_detection = None - for detection in detection_array.detections: - x1, y1, x2, y2 = int(detection.x1), int(detection.y1), int(detection.x2), int(detection.y2) - iou = self.calculate_iou([x1, y1, x2, y2], [x_pred1, y_pred1, x_pred2, y_pred2]) - detection.iou = iou # Set the IoU value - rospy.loginfo(f"Detection {detection.id}: IoU with operator={iou:.2f}") - - # Update the bounding box image with IoU values - label_text = f'#{detection.id} {int(detection.label)}: {detection.score:.2f} IoU={iou:.2f}' - cv2.putText( - bounding_box_image, label_text, (x1, y1 - 10), - cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1 - ) - - # Find the best detection based on IoU - if iou > best_iou: - best_iou = iou - best_detection = detection - - # Update the Kalman Filter with the best detection if IoU is above the threshold - if best_iou > self.iou_threshold and best_detection is not None: - x_center = (best_detection.x1 + best_detection.x2) / 2 - y_center = (best_detection.y1 + best_detection.y2) / 2 - self.kalman_filter_operator.update(np.array([[x_center], [y_center]])) - operator_box = [best_detection.x1, best_detection.y1, best_detection.x2, best_detection.y2] - else: - rospy.logwarn("No detection with IoU above threshold") + cv2.rectangle(bounding_box_image, (x_pred1, y_pred1), (x_pred2, y_pred2), (255, 0, 0), 2) + cv2.circle(bounding_box_image, (int(x_pred), int(y_pred)), 5, (255, 0, 0), -1) + + # Use Nearest Neighbor approach to update Kalman Filter + best_iou = -1 + best_detection = None + for detection in detection_array.detections: + x1, y1, x2, y2 = int(detection.x1), int(detection.y1), int(detection.x2), int(detection.y2) + detection_center = np.array([(x1 + x2) / 2, (y1 + y2) / 2]) + iou = self.calculate_iou([x1, y1, x2, y2], [x_pred1, y_pred1, x_pred2, y_pred2]) + detection.iou = iou # Set the IoU value + rospy.loginfo(f"Detection {detection.id}: IoU with operator={iou:.2f}") + + # Update the bounding box image with IoU values + label_text = f'#{detection.id} {int(detection.label)}: {detection.score:.2f} IoU={iou:.2f}' + cv2.putText( + bounding_box_image, label_text, (x1, y1 - 10), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1 + ) + + if iou > best_iou: + best_iou = iou + best_detection = detection + + # Update the Kalman Filter with the detection with the highest IoU + if best_detection is not None: + x_center = (best_detection.x1 + best_detection.x2) / 2 + y_center = (best_detection.y1 + best_detection.y2) / 2 + self.kalman_filter_operator.update(np.array([[x_center], [y_center]])) + operator_box = [best_detection.x1, best_detection.y1, best_detection.x2, best_detection.y2] + cv2.circle(bounding_box_image, (int(x_center), int(y_center)), 5, (0, 0, 255), -1) + else: + rospy.logwarn("No detection with IoU above threshold, using prediction") + cv2.circle(bounding_box_image, (int(x_pred), int(y_pred)), 5, (255, 0, 0), -1) + cv2.putText( + bounding_box_image, 'Predicted position', (int(x_pred), int(y_pred) - 10), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1 + ) # Publish segmented images as a batch self.segmented_images_pub.publish(segmented_images_msg) @@ -247,7 +262,7 @@ def image_callback(self, data): # Publish bounding box image try: bounding_box_image_msg = self.bridge.cv2_to_imgmsg(bounding_box_image, "bgr8") - bounding_box_image_msg.header.stamp = data.header.stamp # Use the same timestamp + bounding_box_image_msg.header.stamp = data.header.stamp self.bounding_box_image_pub.publish(bounding_box_image_msg) except CvBridgeError as e: rospy.logerr(f"CV Bridge Error while publishing: {e}") @@ -256,19 +271,16 @@ def image_callback(self, data): self.detection_pub.publish(detection_array) def calculate_iou(self, box1, box2): - """Calculate the Intersection over Union (IoU) of two bounding boxes.""" x1_min, y1_min, x1_max, y1_max = box1 x2_min, y2_min, x2_max, y2_max = box2 - # Calculate intersection inter_x_min = max(x1_min, x2_min) inter_y_min = max(y1_min, y2_min) inter_x_max = min(x1_max, x2_max) - inter_y_max = min y1_max) + inter_y_max = min(y1_max, y2_max) inter_area = max(0, inter_x_max - inter_x_min) * max(0, inter_y_max - inter_y_min) - # Calculate union box1_area = (x1_max - x1_min) * (y1_max - y1_min) box2_area = (x2_max - x2_min) * (y2_max - y2_min) union_area = box1_area + box2_area - inter_area