2D->3D detection to pointcloud new pipeline

ed_sensor_integration/CMakeLists.txt

@@ -1,8 +1,29 @@
-cmake_minimum_required(VERSION 3.0.2)
+cmake_minimum_required(VERSION 3.5)
 project(ed_sensor_integration)
 
-add_compile_options(-Wall -Werror=all)
-add_compile_options(-Wextra -Werror=extra)
+
+#add_compile_options(-Wall -Werror=all)
+#add_compile_options(-Wextra -Werror=extra)
+
+# -------------- ONNXRuntime Setup (define this early) ------------------#
+set(ONNXRUNTIME_VERSION 1.21.1)
+set(ONNXRUNTIME_ROOT "/home/amigo/Documents/repos/hero_sam.bak/onnxruntime-linux-x64-gpu-${ONNXRUNTIME_VERSION}")
+
+# -------------- Neural Network Models ------------------#
+
+# Copy model files to the same folder as the executables
+#configure_file(/home/amigo/Documents/repos/hero_sam.bak/yolo_inference/data/coco.yaml ${CATKIN_DEVEL_PREFIX}/lib/ed_sensor_integration/coco.yaml COPYONLY)
+#configure_file(/home/amigo/Documents/repos/hero_sam.bak/yolo_inference/model/yolo11m.onnx ${CATKIN_DEVEL_PREFIX}/lib/ed_sensor_integration/yolo11m.onnx COPYONLY)
+# configure_file(/home/amigo/Documents/repos/hero_sam.bak/sam_inference/model/SAM_mask_decoder.onnx ${CATKIN_DEVEL_PREFIX}/lib/ed_sensor_integration/SAM_mask_decoder.onnx COPYONLY)
+# configure_file(/home/amigo/Documents/repos/hero_sam.bak/sam_inference/model/SAM_encoder.onnx ${CATKIN_DEVEL_PREFIX}/lib/ed_sensor_integration/SAM_encoder.onnx COPYONLY)
+#OR
+# Define model paths
+set(YOLO_MODELS_PATH "/home/amigo/Documents/repos/hero_sam.bak/yolo_inference/model")
+set(SAM_MODELS_PATH "/home/amigo/Documents/repos/hero_sam.bak/sam_inference/model")
+
+# Make these paths available to the C++ code
+add_definitions(-DYOLO_MODELS_PATH="${YOLO_MODELS_PATH}")
+add_definitions(-DSAM_MODELS_PATH="${SAM_MODELS_PATH}")
 
 find_package(OpenCV REQUIRED)
 find_package(PCL REQUIRED COMPONENTS common kdtree)
@@ -21,6 +42,13 @@ find_package(catkin REQUIRED COMPONENTS
   tue_config
   tue_filesystem
   visualization_msgs
+  # 2D -> 3D point cloud estimation
+  yolo_onnx_ros
+  sam_onnx_ros
+  bmm_ros
+  # For displaying SAM MASK
+  cv_bridge
+  image_transport
 )
 
 # ------------------------------------------------------------------------------------------------
@@ -29,8 +57,14 @@ find_package(catkin REQUIRED COMPONENTS
 
 catkin_package(
   INCLUDE_DIRS include
-  LIBRARIES ed_kinect
-  CATKIN_DEPENDS code_profiler ed ${PROJECT_NAME}_msgs geolib2 image_geometry rgbd rgbd_image_buffer roscpp tue_config visualization_msgs
+  LIBRARIES
+    ed_kinect
+    ${PROJECT_NAME}_console_bridge
+    ed_association
+    ed_laser
+  CATKIN_DEPENDS
+    code_profiler ed ${PROJECT_NAME}_msgs geolib2 image_geometry rgbd rgbd_image_buffer roscpp tue_config visualization_msgs
+    sam_onnx_ros bmm_ros cv_bridge image_transport
   DEPENDS OpenCV PCL
 )
 
@@ -43,14 +77,17 @@ include_directories(
   SYSTEM
   ${OpenCV_INCLUDE_DIRS}
   ${catkin_INCLUDE_DIRS}
+  ${ONNXRUNTIME_ROOT}/include
+
 )
 
 # ------------------------------------------------------------------------------------------------
 #                                         LIBRARIES
 # ------------------------------------------------------------------------------------------------
 
+
+
 add_library(${PROJECT_NAME}_console_bridge
-  # Foward geolib2 and tue_filesystem logging to rosconsole
   src/rosconsole_bridge.cpp
 )
 
@@ -60,24 +97,34 @@ add_library(ed_association
 )
 
 add_library(ed_kinect
-    include/ed/kinect/association.h
-    include/ed/kinect/beam_model.h
-    include/ed/kinect/entity_update.h
-    include/ed/kinect/fitter.h
-    include/ed/kinect/mesh_tools.h
-    include/ed/kinect/renderer.h
-    include/ed/kinect/segmenter.h
-    include/ed/kinect/updater.h
-
-    src/kinect/association.cpp
-    src/kinect/beam_model.cpp
-    src/kinect/fitter.cpp
-    src/kinect/mesh_tools.cpp
-    src/kinect/renderer.cpp
-    src/kinect/segmenter.cpp
-    src/kinect/updater.cpp
-)
-target_link_libraries(ed_kinect ${PROJECT_NAME}_console_bridge ed_association ${OpenCV_LIBRARIES} ${catkin_LIBRARIES})
+  include/ed/kinect/association.h
+  include/ed/kinect/beam_model.h
+  include/ed/kinect/entity_update.h
+  include/ed/kinect/fitter.h
+  include/ed/kinect/mesh_tools.h
+  include/ed/kinect/renderer.h
+  include/ed/kinect/segmenter.h
+  include/ed/kinect/updater.h
+  src/kinect/association.cpp
+  src/kinect/beam_model.cpp
+  src/kinect/fitter.cpp
+  src/kinect/mesh_tools.cpp
+  src/kinect/renderer.cpp
+  src/kinect/segmenter.cpp
+  src/kinect/updater.cpp
+)
+
+add_library(ed_sam_segment_module
+  include/ed_sensor_integration/kinect/segmodules/sam_seg_module.h
+  src/kinect/sam_seg_module.cpp
+)
+target_link_libraries(ed_kinect
+  ${PROJECT_NAME}_console_bridge
+  ed_association
+  ed_sam_segment_module
+  ${OpenCV_LIBRARIES}
+  ${catkin_LIBRARIES}
+)
 add_dependencies(ed_kinect ${catkin_EXPORTED_TARGETS})
 
 add_library(ed_laser
@@ -93,28 +140,36 @@ add_dependencies(ed_laser ${catkin_EXPORTED_TARGETS})
 # ------------------------------------------------------------------------------------------------
 
 add_library(ed_kinect_plugin
-    src/kinect/kinect_plugin.cpp
-    src/kinect/kinect_plugin.h
-    src/kinect/ray_tracer.cpp
-    src/kinect/ray_tracer.h
+  src/kinect/kinect_plugin.cpp
+  src/kinect/kinect_plugin.h
+  src/kinect/ray_tracer.cpp
+  src/kinect/ray_tracer.h
 )
-target_link_libraries(ed_kinect_plugin ${PROJECT_NAME}_console_bridge ed_kinect ${catkin_LIBRARIES})
+
+target_link_libraries(ed_kinect_plugin
+  ${PROJECT_NAME}_console_bridge
+  ed_kinect
+  ${catkin_LIBRARIES}
+)
+
 add_dependencies(ed_kinect_plugin ${catkin_EXPORTED_TARGETS})
 
 # ------------------------------------------------------------------------------------------------
 
 add_library(ed_laser_plugin
-    src/laser/laser_plugin.cpp
-    src/laser/laser_plugin.h
+  src/laser/laser_plugin.cpp
+  src/laser/laser_plugin.h
 )
+
 target_link_libraries(ed_laser_plugin ${PROJECT_NAME}_console_bridge ed_laser ed_association ${catkin_LIBRARIES})
 
 # ------------------------------------------------------------------------------------------------
 
 add_library(ed_clearer_plugin
-    src/clearer/clearer_plugin.cpp
-    src/clearer/clearer_plugin.h
+  src/clearer/clearer_plugin.cpp
+  src/clearer/clearer_plugin.h
 )
+
 target_link_libraries(ed_clearer_plugin ${catkin_LIBRARIES})
 
 # ------------------------------------------------------------------------------------------------
@@ -126,13 +181,26 @@ target_link_libraries(ed_image_saver ${catkin_LIBRARIES})
 add_dependencies(ed_image_saver ${catkin_EXPORTED_TARGETS})
 
 add_executable(ed_segmenter tools/segmenter.cpp)
-target_link_libraries(ed_segmenter ed_kinect ${catkin_LIBRARIES})
+target_link_libraries(ed_segmenter
+  ed_kinect
+  ${catkin_LIBRARIES}
+)
 
 add_executable(ed_fitter_data tools/fitter_viz_data.cpp)
-target_link_libraries(ed_fitter_data ${PROJECT_NAME}_console_bridge ed_kinect ${OpenCV_LIBRARIES} ${catkin_LIBRARIES})
+target_link_libraries(ed_fitter_data
+  ${PROJECT_NAME}_console_bridge
+  ed_kinect
+  ${OpenCV_LIBRARIES}
+  ${catkin_LIBRARIES}
+)
 
 add_executable(ed_fitter_live tools/fitter_viz_live.cpp)
-target_link_libraries(ed_fitter_live ${PROJECT_NAME}_console_bridge ed_kinect ${OpenCV_LIBRARIES} ${catkin_LIBRARIES})
+target_link_libraries(ed_fitter_live
+  ${PROJECT_NAME}_console_bridge
+  ed_kinect
+  ${OpenCV_LIBRARIES}
+  ${catkin_LIBRARIES}
+)
 
 # ------------------------------------------------------------------------------------------------
 #                                          Install
@@ -179,11 +247,18 @@ if (CATKIN_ENABLE_TESTING)
   catkin_add_catkin_lint_test("-W2 --ignore HEADER_OUTSIDE_PACKAGE_INCLUDE_PATH")
 
   catkin_add_gtest(test_furniture_fitting test/test_furniture_fit.cpp)
-  target_link_libraries(test_furniture_fitting ${PROJECT_NAME}_console_bridge ed_kinect ${OpenCV_LIBRARIES} ${catkin_LIBRARIES})
+  target_link_libraries(test_furniture_fitting
+    ${PROJECT_NAME}_console_bridge
+    ed_kinect
+    ${OpenCV_LIBRARIES}
+    ${catkin_LIBRARIES}
+  )
 
   catkin_add_gtest(test_laser_fitting test/test_laser_segmenter.cpp)
-  target_link_libraries(test_laser_fitting ${PROJECT_NAME}_console_bridge ed_laser ${OpenCV_LIBRARIES} ${catkin_LIBRARIES})
+  target_link_libraries(test_laser_fitting
+    ${PROJECT_NAME}_console_bridge
+    ed_laser
+    ${OpenCV_LIBRARIES}
+    ${catkin_LIBRARIES}
+  )
 endif ()
-
-
-

ed_sensor_integration/include/ed/kinect/segmenter.h

@@ -6,6 +6,7 @@
 #include <rgbd/types.h>
 #include <geolib/datatypes.h>
 #include <opencv2/core/core.hpp>
+#include <tue/config/configuration.h>
 
 #include <ed/convex_hull.h>
 #include <ed/types.h>
@@ -29,7 +30,7 @@ class Segmenter
 
 public:
 
-    Segmenter();
+    Segmenter(tue::Configuration config);
 
     ~Segmenter();
 
@@ -38,12 +39,31 @@ class Segmenter
 
     void calculatePointsWithin(const rgbd::Image& image, const geo::Shape& shape,
                                const geo::Pose3D& shape_pose, cv::Mat& filtered_depth_image) const;
-
-    void cluster(const cv::Mat& depth_image, const geo::DepthCamera& cam_model,
-                 const geo::Pose3D& sensor_pose, std::vector<EntityUpdate>& clusters) const;
+    /**
+     * @brief Preprocess RGB image for segmentation. This function masks the RGB image with the filtered depth image (non-zero depth values are kept)
+     * It returns a masked RGB image with values only where depth is non-zero
+     *
+     * @param rgb_image
+     * @param filtered_depth_image
+     * @return cv::Mat filtered_depth_image
+     */
+    cv::Mat preprocessRGBForSegmentation(const cv::Mat& rgb_image, const cv::Mat& filtered_depth_image) const;
+
+    /**
+     * @brief Cluster the depth image into segments. Applies new algorithm which is the YOLO - SAM - BMM depth segmentation pipeline.
+     *
+     * @param depth_image
+     * @param cam_model
+     * @param sensor_pose
+     * @param clusters
+     * @param rgb_image
+     * @return std::vector<cv::Mat> masks // 3D pointcloud masks of all the segmented objects
+     */
+    std::vector<cv::Mat> cluster(const cv::Mat& depth_image, const geo::DepthCamera& cam_model,
+                 const geo::Pose3D& sensor_pose, std::vector<EntityUpdate>& clusters, const cv::Mat& rgb_image);
 
 private:
-
+    tue::Configuration config_;
 };
 
 #endif

ed_sensor_integration/include/ed/kinect/updater.h

@@ -5,9 +5,10 @@
 #include "ed/kinect/segmenter.h"
 #include "ed/kinect/entity_update.h"
 
+#include <tue/config/configuration.h>
+
 #include <map>
 #include <vector>
-
 // ----------------------------------------------------------------------------------------------------
 
 struct UpdateRequest
@@ -48,7 +49,7 @@ class Updater
 
 public:
 
-    Updater();
+    Updater(tue::Configuration config);
 
     ~Updater();
 
@@ -59,11 +60,15 @@ class Updater
 
     Fitter fitter_;
 
-    Segmenter segmenter_;
+    std::unique_ptr<Segmenter> segmenter_;
 
     // Stores for each segmented entity with which area description it was found
     std::map<ed::UUID, std::string> id_to_area_description_;
 
+    //For displaying SAM MASK
+    ros::Publisher mask_pub_;
+    ros::Publisher cloud_pub_;
+
 };
 
 #endif

ed_sensor_integration/include/ed_sensor_integration/kinect/segmodules/sam_seg_module.h

@@ -0,0 +1,53 @@
+#pragma once
+
+#include <ros/ros.h>
+#include "yolo_inference.h"
+#include "sam_inference.h"
+#include <sensor_msgs/Image.h>
+#include <sensor_msgs/PointCloud2.h>
+#include <opencv2/core.hpp>
+//#include "ed/kinect/entity_update.h"
+#include <cv_bridge/cv_bridge.h>
+#include <geolib/sensors/DepthCamera.h>
+#include <geolib/Shape.h>
+#include <pcl/point_types.h>
+#include <pcl/point_cloud.h>
+#include <pcl/segmentation/extract_clusters.h>
+#include <pcl_conversions/pcl_conversions.h>
+#include <sensor_msgs/Image.h>
+#include <sensor_msgs/PointCloud2.h>
+
+// EntityUpdate and UpdateResult come from kinect (avoid including laser variant to prevent redefinition)
+#include <ed/kinect/entity_update.h>
+#include <ed/kinect/segmenter.h>  // defines UpdateResult (and possibly other needed types)
+
+/**
+ * @brief Segmentation pipeline that processes the input image and generates segmentation masks.
+ *
+ * @param img The input RGB image to segment.
+ * @return std::vector<cv::Mat> The generated segmentation masks.
+ */
+std::vector<cv::Mat> SegmentationPipeline(const cv::Mat& img);
+
+/**
+ * @brief Overlay segmentation masks on the RGB image for visualization purposes.
+ *
+ * @param rgb The RGB image to overlay masks on.
+ * @param masks The segmentation masks to overlay.
+ */
+void overlayMasksOnImage_(cv::Mat& rgb, const std::vector<cv::Mat>& masks);
+
+/**
+ * @brief Publish segmentation results and pointcloud estimation as ROS messages.
+ *
+ * @param filtered_depth_image The filtered depth image to publish.
+ * @param rgb The RGB image to publish.
+ * @param sensor_pose The pose of the sensor.
+ * @param clustered_images The clustered segmentation masks.
+ * @param mask_pub_ The ROS publisher for the mask images.
+ * @param cloud_pub_ The ROS publisher for the point cloud data.
+ * @param res_updates The entity updates to publish.
+ */
+void publishSegmentationResults(const cv::Mat& filtered_depth_image, const cv::Mat& rgb,
+                                const geo::Pose3D& sensor_pose, std::vector<cv::Mat>& clustered_images,
+                                ros::Publisher& mask_pub_, ros::Publisher& cloud_pub_, std::vector<EntityUpdate>& res_updates);