add support for processing points with masks

CMakeLists.txt

@@ -7,6 +7,7 @@ if(NOT DEFINED CMAKE_SUPPRESS_DEVELOPER_WARNINGS)
 endif()
 
 find_package(catkin REQUIRED COMPONENTS
+  cv_bridge
   ros_numpy
   geometry_msgs
   image_geometry
@@ -43,6 +44,7 @@ generate_messages(
 catkin_package(
   INCLUDE_DIRS include
   CATKIN_DEPENDS
+    cv_bridge
     ros_numpy
     geometry_msgs
     image_geometry

include/predict_with_cloud_node/predict_with_cloud_node.h

@@ -1,73 +1,87 @@
 #pragma once
 
+#include <ros/ros.h>
+#include <sensor_msgs/CameraInfo.h>
+#include <sensor_msgs/PointCloud2.h>
+#include <std_msgs/Header.h>
+#include <visualization_msgs/Marker.h>
+#include <visualization_msgs/MarkerArray.h>
+#include <vision_msgs/Detection2DArray.h>
+#include <vision_msgs/Detection3D.h>
+#include <vision_msgs/Detection3DArray.h>
+#include <ultralytics_ros/YoloResult.h>
 #include <geometry_msgs/TransformStamped.h>
-#include <image_geometry/pinhole_camera_model.h>
+#include <tf2_ros/transform_listener.h>
 #include <message_filters/subscriber.h>
 #include <message_filters/sync_policies/approximate_time.h>
 #include <message_filters/synchronizer.h>
+#include <image_geometry/pinhole_camera_model.h>
+
 #include <pcl/common/common.h>
 #include <pcl/common/pca.h>
 #include <pcl/common/transforms.h>
 #include <pcl/io/pcd_io.h>
 #include <pcl/point_types.h>
 #include <pcl/segmentation/extract_clusters.h>
+#include <pcl/filters/voxel_grid.h>
 #include <pcl_conversions/pcl_conversions.h>
 #include <pcl_ros/point_cloud.h>
 #include <pcl_ros/transforms.h>
-#include <ros/ros.h>
-#include <sensor_msgs/CameraInfo.h>
-#include <sensor_msgs/PointCloud2.h>
-#include <std_msgs/Header.h>
-#include <tf2_ros/transform_listener.h>
-#include <vision_msgs/Detection2DArray.h>
-#include <vision_msgs/Detection3D.h>
-#include <vision_msgs/Detection3DArray.h>
-#include <vision_msgs/ObjectHypothesisWithPose.h>
-#include <visualization_msgs/Marker.h>
-#include <visualization_msgs/MarkerArray.h>
 
-#include <Eigen/Dense>
 #include <opencv2/opencv.hpp>
+#include <cv_bridge/cv_bridge.h>
+#include <Eigen/Dense>
+
+typedef message_filters::sync_policies::ApproximateTime<sensor_msgs::CameraInfo, sensor_msgs::PointCloud2,
+                                                        ultralytics_ros::YoloResult>
+    ApproximateSyncPolicy;
 
 class PredictWithCloudNode
 {
 private:
-  ros::NodeHandle _nh;
-  ros::NodeHandle _pnh;
-  std::string _camera_info_topic;
-  std::string _lidar_topic;
-  std::string _detection2d_topic;
-  std::string _detection3d_topic;
-  ros::Publisher _detection_cloud_pub;
-  ros::Publisher _detection3d_pub;
-  ros::Publisher _marker_pub;
-  typedef message_filters::sync_policies::ApproximateTime<sensor_msgs::CameraInfo, sensor_msgs::PointCloud2,
-                                                          vision_msgs::Detection2DArray>
-      _sensor_fusion_sync_subs;
-  message_filters::Subscriber<sensor_msgs::CameraInfo> _camera_info_sub;
-  message_filters::Subscriber<sensor_msgs::PointCloud2> _lidar_sub;
-  message_filters::Subscriber<vision_msgs::Detection2DArray> _detection2d_sub;
-  boost::shared_ptr<message_filters::Synchronizer<_sensor_fusion_sync_subs>> _sensor_fusion_sync;
-  boost::shared_ptr<tf2_ros::Buffer> _tf_buffer;
-  boost::shared_ptr<tf2_ros::TransformListener> _tf_listener;
-  image_geometry::PinholeCameraModel _cam_model;
-  float _cluster_tolerance;
-  int _min_cluster_size;
-  int _max_cluster_size;
+  ros::NodeHandle nh_;
+  ros::NodeHandle pnh_;
+  ros::Publisher detection_cloud_pub_;
+  ros::Publisher detection3d_pub_;
+  ros::Publisher marker_pub_;
+  ros::Time last_call_time_;
+  message_filters::Subscriber<sensor_msgs::CameraInfo> camera_info_sub_;
+  message_filters::Subscriber<sensor_msgs::PointCloud2> lidar_sub_;
+  message_filters::Subscriber<ultralytics_ros::YoloResult> yolo_result_sub_;
+  boost::shared_ptr<message_filters::Synchronizer<ApproximateSyncPolicy>> sync_;
+  boost::shared_ptr<tf2_ros::Buffer> tf_buffer_;
+  boost::shared_ptr<tf2_ros::TransformListener> tf_listener_;
+  image_geometry::PinholeCameraModel cam_model_;
+  std::string camera_info_topic_;
+  std::string lidar_topic_;
+  std::string yolo_result_topic_;
+  std::string yolo_3d_result_topic_;
+  float cluster_tolerance_;
+  float voxel_leaf_size_;
+  int min_cluster_size_;
+  int max_cluster_size_;
 
 public:
   PredictWithCloudNode();
   void syncCallback(const sensor_msgs::CameraInfo::ConstPtr& camera_info_msg,
                     const sensor_msgs::PointCloud2ConstPtr& cloud_msg,
-                    const vision_msgs::Detection2DArrayConstPtr& detections2d_msg);
-  pcl::PointCloud<pcl::PointXYZ> msg2TransformedCloud(const sensor_msgs::PointCloud2ConstPtr& cloud_msg);
-  std::tuple<vision_msgs::Detection3DArray, sensor_msgs::PointCloud2>
-  projectCloud(const pcl::PointCloud<pcl::PointXYZ>& cloud,
-               const vision_msgs::Detection2DArrayConstPtr& detections2d_msg, const std_msgs::Header& header);
-  pcl::PointCloud<pcl::PointXYZ> cloud2TransformedCloud(const pcl::PointCloud<pcl::PointXYZ>& cloud,
-                                                        const std_msgs::Header& header);
-  pcl::PointCloud<pcl::PointXYZ> euclideanClusterExtraction(const pcl::PointCloud<pcl::PointXYZ>& cloud);
-  void createBoundingBox(vision_msgs::Detection3DArray& detections3d_msg, const pcl::PointCloud<pcl::PointXYZ>& cloud,
+                    const ultralytics_ros::YoloResult::ConstPtr& yolo_result_msg);
+  void projectCloud(const pcl::PointCloud<pcl::PointXYZ>::Ptr& cloud,
+                    const ultralytics_ros::YoloResultConstPtr& yolo_result_msg, const std_msgs::Header& header,
+                    vision_msgs::Detection3DArray& detections3d_msg,
+                    sensor_msgs::PointCloud2& combine_detection_cloud_msg);
+  void processPointsWithBbox(const pcl::PointCloud<pcl::PointXYZ>::Ptr& cloud,
+                             const vision_msgs::Detection2D& detection,
+                             pcl::PointCloud<pcl::PointXYZ>::Ptr& detection_cloud_raw);
+  void processPointsWithMask(const pcl::PointCloud<pcl::PointXYZ>::Ptr& cloud, const sensor_msgs::Image& mask,
+                             pcl::PointCloud<pcl::PointXYZ>::Ptr& detection_cloud_raw);
+  pcl::PointCloud<pcl::PointXYZ>::Ptr msg2TransformedCloud(const sensor_msgs::PointCloud2ConstPtr& cloud_msg);
+  pcl::PointCloud<pcl::PointXYZ>::Ptr cloud2TransformedCloud(const pcl::PointCloud<pcl::PointXYZ>::Ptr& cloud,
+                                                             const std_msgs::Header& header);
+  pcl::PointCloud<pcl::PointXYZ>::Ptr euclideanClusterExtraction(const pcl::PointCloud<pcl::PointXYZ>::Ptr& cloud);
+  void createBoundingBox(vision_msgs::Detection3DArray& detections3d_msg,
+                         const pcl::PointCloud<pcl::PointXYZ>::Ptr& cloud,
                          const std::vector<vision_msgs::ObjectHypothesisWithPose>& detections_results);
-  visualization_msgs::MarkerArray createMarkerArray(const vision_msgs::Detection3DArray& detections3d_msg);
+  visualization_msgs::MarkerArray createMarkerArray(const vision_msgs::Detection3DArray& detections3d_msg,
+                                                    const double& duration);
 };

launch/kitti_predict_with_cloud.launch

@@ -1,59 +1,58 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <launch>
-  <arg name="yolo_model" default="yolov8n.pt"/>
-  <arg name="detection_topic" default="detection_result"/>
-  <arg name="image_topic" default="kitti/camera_color_left/image_raw"/>
+  <arg name="debug" default="true"/>
+  <arg name="yolo_model" default="yolov8m-seg.pt"/>
+  <arg name="input_topic" default="/kitti/camera_color_left/image_raw"/>
+  <arg name="result_topic" default="/yolo_result"/>
+  <arg name="result_image_topic" default="/yolo_image"/>
   <arg name="conf_thres" default="0.25"/>
   <arg name="iou_thres" default="0.45"/>
   <arg name="max_det" default="300"/>
-  <arg name="debug" default="true"/>
-  <arg name="debug_conf" default="true"/>
-  <arg name="debug_line_width" default="1"/>
-  <arg name="debug_font_size" default="1"/>
-  <arg name="debug_font" default="Arial.ttf"/>
-  <arg name="debug_labels" default="true"/>
-  <arg name="debug_boxes" default="true"/>
-  <arg name="camera_info_topic" default="kitti/camera_color_left/camera_info"/>
-  <arg name="lidar_topic" default="kitti/velo/pointcloud"/>
-  <arg name="detection3d_topic" default="detection3d_result"/>
-  <arg name="cluster_tolerance" default="0.5"/>
+  <arg name="classes" default=""/>
+  <arg name="device" default=""/>
+  <arg name="result_conf" default="true"/>
+  <arg name="result_line_width" default="1"/>
+  <arg name="result_font_size" default="1"/>
+  <arg name="result_font" default="Arial.ttf"/>
+  <arg name="result_labels" default="true"/>
+  <arg name="result_boxes" default="true"/>
+
+  <arg name="camera_info_topic" default="/kitti/camera_color_left/camera_info"/>
+  <arg name="lidar_topic" default="/kitti/velo/pointcloud"/>
+  <arg name="yolo_3d_result_topic" default="/yolo_3d_result"/>
+  <arg name="cluster_tolerance" default="0.3"/>
+  <arg name="voxel_leaf_size" default="0.1"/>
   <arg name="min_cluster_size" default="100"/>
-  <arg name="max_cluster_size" default="25000"/>
-  <arg name="use_sim_time" default="true"/>
-  <param if="$(arg use_sim_time)" name="use_sim_time" value="true"/>
-  <!-- Object detection node -->
+  <arg name="max_cluster_size" default="10000"/>
+
   <node name="predict_node" pkg="ultralytics_ros" type="predict_node.py" output="screen">
     <param name="yolo_model" value="$(arg yolo_model)"/>
-    <param name="detection_topic" value="$(arg detection_topic)"/>
-    <param name="image_topic" value="$(arg image_topic)"/>
+    <param name="input_topic" value="$(arg input_topic)"/>
+    <param name="result_topic" value="$(arg result_topic)"/>
+    <param name="result_image_topic" value="$(arg result_image_topic)"/>
     <param name="conf_thres" value="$(arg conf_thres)"/>
     <param name="iou_thres" value="$(arg iou_thres)"/>
     <param name="max_det" value="$(arg max_det)"/>
-    <param name="debug" value="$(arg debug)"/>
-    <param name="debug_conf" value="$(arg debug_conf)"/>
-    <param name="debug_line_width" value="$(arg debug_line_width)"/>
-    <param name="debug_font_size" value="$(arg debug_font_size)"/>
-    <param name="debug_font" value="$(arg debug_font)"/>
-    <param name="debug_labels" value="$(arg debug_labels)"/>
-    <param name="debug_boxes" value="$(arg debug_boxes)"/>
-    <rosparam>
-      # classes: [0] (person)
-      # classes: [1] (bicycle)
-      # classes: [2] (car)
-      # ...
-      # See also : https://github.com/ultralytics/ultralytics/blob/main/ultralytics/datasets/coco128.yaml
-    </rosparam>
+    <param name="result_conf" value="$(arg result_conf)"/>
+    <param name="result_line_width" value="$(arg result_line_width)"/>
+    <param name="result_font_size" value="$(arg result_font_size)"/>
+    <param name="result_font" value="$(arg result_font)"/>
+    <param name="result_labels" value="$(arg result_labels)"/>
+    <param name="result_boxes" value="$(arg result_boxes)"/>
+    <rosparam param="classes" subst_value="true">$(arg classes)</rosparam>
+    <rosparam param="device" subst_value="true">$(arg device)</rosparam>
   </node>
-  <!-- Object detection with pointcloud node -->
+
   <node name="predict_with_cloud_node" pkg="ultralytics_ros" type="predict_with_cloud_node" output="screen">
     <param name="camera_info_topic" value="$(arg camera_info_topic)"/>
     <param name="lidar_topic" value="$(arg lidar_topic)"/>
-    <param name="detection2d_topic" value="$(arg detection_topic)"/>
-    <param name="detection3d_topic" value="$(arg detection3d_topic)"/>
+    <param name="yolo_result_topic" value="$(arg result_topic)"/>
+    <param name="yolo_3d_result_topic" value="$(arg yolo_3d_result_topic)"/>
     <param name="cluster_tolerance" value="$(arg cluster_tolerance)"/>
+    <param name="voxel_leaf_size" value="$(arg voxel_leaf_size)"/>
     <param name="min_cluster_size" value="$(arg min_cluster_size)"/>
     <param name="max_cluster_size" value="$(arg max_cluster_size)"/>
   </node>
-  <!-- rviz -->
+
   <node if="$(arg debug)" pkg="rviz" type="rviz" name="rviz" args="-d $(find ultralytics_ros)/rviz/kitti.rviz"/>
 </launch>

launch/predict_with_cloud.launch

@@ -1,59 +1,58 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <launch>
-  <arg name="yolo_model" default="yolov8n.pt"/>
-  <arg name="detection_topic" default="detection_result"/>
-  <arg name="image_topic" default="image_raw"/>
+  <arg name="debug" default="false"/>
+  <arg name="yolo_model" default="yolov8m-seg.pt"/>
+  <arg name="input_topic" default="/image_raw"/>
+  <arg name="result_topic" default="/yolo_result"/>
+  <arg name="result_image_topic" default="/yolo_image"/>
   <arg name="conf_thres" default="0.25"/>
   <arg name="iou_thres" default="0.45"/>
   <arg name="max_det" default="300"/>
-  <arg name="debug" default="false"/>
-  <arg name="debug_conf" default="true"/>
-  <arg name="debug_line_width" default="1"/>
-  <arg name="debug_font_size" default="1"/>
-  <arg name="debug_font" default="Arial.ttf"/>
-  <arg name="debug_labels" default="true"/>
-  <arg name="debug_boxes" default="true"/>
-  <arg name="camera_info_topic" default="camera_info"/>
-  <arg name="lidar_topic" default="points_raw"/>
-  <arg name="detection3d_topic" default="detection3d_result"/>
-  <arg name="cluster_tolerance" default="0.5"/>
+  <arg name="classes" default=""/>
+  <arg name="device" default=""/>
+  <arg name="result_conf" default="true"/>
+  <arg name="result_line_width" default="1"/>
+  <arg name="result_font_size" default="1"/>
+  <arg name="result_font" default="Arial.ttf"/>
+  <arg name="result_labels" default="true"/>
+  <arg name="result_boxes" default="true"/>
+
+  <arg name="camera_info_topic" default="/camera_info"/>
+  <arg name="lidar_topic" default="/points_raw"/>
+  <arg name="yolo_3d_result_topic" default="/yolo_3d_result"/>
+  <arg name="cluster_tolerance" default="0.3"/>
+  <arg name="voxel_leaf_size" default="0.1"/>
   <arg name="min_cluster_size" default="100"/>
-  <arg name="max_cluster_size" default="25000"/>
-  <arg name="use_sim_time" default="true"/>
-  <param if="$(arg use_sim_time)" name="use_sim_time" value="true"/>
-  <!-- Object detection node -->
+  <arg name="max_cluster_size" default="10000"/>
+
   <node name="predict_node" pkg="ultralytics_ros" type="predict_node.py" output="screen">
     <param name="yolo_model" value="$(arg yolo_model)"/>
-    <param name="detection_topic" value="$(arg detection_topic)"/>
-    <param name="image_topic" value="$(arg image_topic)"/>
+    <param name="input_topic" value="$(arg input_topic)"/>
+    <param name="result_topic" value="$(arg result_topic)"/>
+    <param name="result_image_topic" value="$(arg result_image_topic)"/>
     <param name="conf_thres" value="$(arg conf_thres)"/>
     <param name="iou_thres" value="$(arg iou_thres)"/>
     <param name="max_det" value="$(arg max_det)"/>
-    <param name="debug" value="$(arg debug)"/>
-    <param name="debug_conf" value="$(arg debug_conf)"/>
-    <param name="debug_line_width" value="$(arg debug_line_width)"/>
-    <param name="debug_font_size" value="$(arg debug_font_size)"/>
-    <param name="debug_font" value="$(arg debug_font)"/>
-    <param name="debug_labels" value="$(arg debug_labels)"/>
-    <param name="debug_boxes" value="$(arg debug_boxes)"/>
-    <rosparam>
-      # classes: [0] (person)
-      # classes: [1] (bicycle)
-      # classes: [2] (car)
-      # ...
-      # See also : https://github.com/ultralytics/ultralytics/blob/main/ultralytics/datasets/coco128.yaml
-    </rosparam>
+    <param name="result_conf" value="$(arg result_conf)"/>
+    <param name="result_line_width" value="$(arg result_line_width)"/>
+    <param name="result_font_size" value="$(arg result_font_size)"/>
+    <param name="result_font" value="$(arg result_font)"/>
+    <param name="result_labels" value="$(arg result_labels)"/>
+    <param name="result_boxes" value="$(arg result_boxes)"/>
+    <rosparam param="classes" subst_value="true">$(arg classes)</rosparam>
+    <rosparam param="device" subst_value="true">$(arg device)</rosparam>
   </node>
-  <!-- Object detection with pointcloud node -->
+
   <node name="predict_with_cloud_node" pkg="ultralytics_ros" type="predict_with_cloud_node" output="screen">
     <param name="camera_info_topic" value="$(arg camera_info_topic)"/>
     <param name="lidar_topic" value="$(arg lidar_topic)"/>
-    <param name="detection2d_topic" value="$(arg detection_topic)"/>
-    <param name="detection3d_topic" value="$(arg detection3d_topic)"/>
+    <param name="yolo_result_topic" value="$(arg result_topic)"/>
+    <param name="yolo_3d_result_topic" value="$(arg yolo_3d_result_topic)"/>
     <param name="cluster_tolerance" value="$(arg cluster_tolerance)"/>
+    <param name="voxel_leaf_size" value="$(arg voxel_leaf_size)"/>
     <param name="min_cluster_size" value="$(arg min_cluster_size)"/>
     <param name="max_cluster_size" value="$(arg max_cluster_size)"/>
   </node>
-  <!-- rviz -->
+
   <node if="$(arg debug)" pkg="rviz" type="rviz" name="rviz" args="-d $(find ultralytics_ros)/rviz/default.rviz"/>
 </launch>

package.xml

@@ -6,6 +6,7 @@
   <maintainer email="[email protected]">Alpaca-zip</maintainer>
   <license>TODO</license>
   <buildtool_depend>catkin</buildtool_depend>
+  <depend>cv_bridge</depend>
   <depend>ros_numpy</depend>
   <depend>geometry_msgs</depend>
   <depend>image_geometry</depend>

rviz/default.rviz

@@ -97,7 +97,7 @@ Visualization Manager:
       Value: true
     - Class: rviz/Image
       Enabled: true
-      Image Topic: /debug_image
+      Image Topic: /yolo_image
       Max Value: 1
       Median window: 5
       Min Value: 0

rviz/kitti.rviz

@@ -98,7 +98,7 @@ Visualization Manager:
       Value: true
     - Class: rviz/Image
       Enabled: true
-      Image Topic: /debug_image
+      Image Topic: /yolo_image
       Max Value: 1
       Median window: 5
       Min Value: 0