diff --git a/CMakeLists.txt b/CMakeLists.txt index 330d97eb..badb7105 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -31,6 +31,8 @@ find_package(PCL REQUIRED) add_message_files( FILES YoloResult.msg + Detection2D.msg + Detection2DArray.msg ) generate_messages( diff --git a/include/tracker_with_cloud_node/tracker_with_cloud_node.h b/include/tracker_with_cloud_node/tracker_with_cloud_node.h index 6b98ae5a..c9709819 100644 --- a/include/tracker_with_cloud_node/tracker_with_cloud_node.h +++ b/include/tracker_with_cloud_node/tracker_with_cloud_node.h @@ -24,10 +24,12 @@ #include #include #include -#include #include #include + #include +#include + #include #include #include @@ -89,7 +91,7 @@ class TrackerWithCloudNode vision_msgs::Detection3DArray& detections3d_msg, sensor_msgs::PointCloud2& combine_detection_cloud_msg); void processPointsWithBbox(const pcl::PointCloud::Ptr& cloud, - const vision_msgs::Detection2D& detection, + const ultralytics_ros::Detection2D& detection, pcl::PointCloud::Ptr& detection_cloud_raw); void processPointsWithMask(const pcl::PointCloud::Ptr& cloud, const sensor_msgs::Image& mask, pcl::PointCloud::Ptr& detection_cloud_raw); diff --git a/msg/Detection2D.msg b/msg/Detection2D.msg new file mode 100644 index 00000000..ededf20b --- /dev/null +++ b/msg/Detection2D.msg @@ -0,0 +1,17 @@ +# Defines a 2D detection result. +# +# This is similar to a 2D classification, but includes position information, +# allowing a classification result for a specific crop or image point to +# to be located in the larger image. + +std_msgs/Header header +int8 track_id +# Class probabilities +vision_msgs/ObjectHypothesisWithPose[] results + +# 2D bounding box surrounding the object. +vision_msgs/BoundingBox2D bbox + +# The 2D data that generated these results (i.e. region proposal cropped out of +# the image). Not required for all use cases, so it may be empty. +sensor_msgs/Image source_img \ No newline at end of file diff --git a/msg/Detection2DArray.msg b/msg/Detection2DArray.msg new file mode 100644 index 00000000..7d05be04 --- /dev/null +++ b/msg/Detection2DArray.msg @@ -0,0 +1,7 @@ +# A list of 2D detections, for a multi-object 2D detector. + +std_msgs/Header header + +# A list of the detected proposals. A multi-proposal detector might generate +# this list with many candidate detections generated from a single input. +Detection2D[] detections diff --git a/msg/YoloResult.msg b/msg/YoloResult.msg index c00e0b1a..1857fb7e 100644 --- a/msg/YoloResult.msg +++ b/msg/YoloResult.msg @@ -1,3 +1,3 @@ std_msgs/Header header -vision_msgs/Detection2DArray detections -sensor_msgs/Image[] masks +Detection2DArray detections +sensor_msgs/Image[] masks \ No newline at end of file diff --git a/script/tracker_node.py b/script/tracker_node.py index 2d78145d..dd74a3c6 100755 --- a/script/tracker_node.py +++ b/script/tracker_node.py @@ -23,8 +23,8 @@ import rospy from sensor_msgs.msg import Image from ultralytics import YOLO -from vision_msgs.msg import Detection2D, Detection2DArray, ObjectHypothesisWithPose -from ultralytics_ros.msg import YoloResult +from vision_msgs.msg import ObjectHypothesisWithPose +from ultralytics_ros.msg import YoloResult, Detection2D, Detection2DArray class TrackerNode: @@ -91,11 +91,15 @@ def image_callback(self, msg): def create_detections_array(self, results): detections_msg = Detection2DArray() + track_ids = results[0].boxes.id bounding_box = results[0].boxes.xywh classes = results[0].boxes.cls confidence_score = results[0].boxes.conf - for bbox, cls, conf in zip(bounding_box, classes, confidence_score): + if track_ids is None: + track_ids = [0] + for bbox, cls, conf, track_id in zip(bounding_box, classes, confidence_score, track_ids): detection = Detection2D() + detection.track_id = int(track_id) detection.bbox.center.x = float(bbox[0]) detection.bbox.center.y = float(bbox[1]) detection.bbox.size_x = float(bbox[2]) diff --git a/src/tracker_with_cloud_node.cpp b/src/tracker_with_cloud_node.cpp index c8f56687..c8e0ce08 100644 --- a/src/tracker_with_cloud_node.cpp +++ b/src/tracker_with_cloud_node.cpp @@ -106,7 +106,7 @@ void TrackerWithCloudNode::projectCloud(const pcl::PointCloud::Pt } void TrackerWithCloudNode::processPointsWithBbox(const pcl::PointCloud::Ptr& cloud, - const vision_msgs::Detection2D& detection, + const ultralytics_ros::Detection2D& detection, pcl::PointCloud::Ptr& detection_cloud_raw) { for (const auto& point : cloud->points)