feat: add IPM projection

Signed-off-by: kminoda <koji.minoda@tier4.jp>

CMakeLists.txt

@@ -34,6 +34,8 @@ if(BUILD_ROS2_LIGHTNET_TRT)
   find_package(autoware_cmake REQUIRED)
   autoware_package()
 
+  find_package(Eigen3 REQUIRED)
+
   # build nodelet
   # TODO
 
@@ -42,11 +44,12 @@ if(BUILD_ROS2_LIGHTNET_TRT)
     src/main.cpp
     src/lightnet_trt_node.cpp
     src/lightnet_trt_core.cpp
+    src/utils.cpp
   )
   ament_target_dependencies(lightnet_trt_node)
   target_include_directories(lightnet_trt_node PRIVATE modules/ extra/)
   target_link_libraries(lightnet_trt_node detector)
-  target_link_libraries(lightnet_trt_node ${CUDA_LIBRARIES} ${OpenCV_LIBS} nvinfer nvinfer_plugin nvparsers gflags boost_system boost_filesystem "stdc++fs" "stdc++")
+  target_link_libraries(lightnet_trt_node ${CUDA_LIBRARIES} ${OpenCV_LIBS} Eigen3::Eigen nvinfer nvinfer_plugin nvparsers gflags boost_system boost_filesystem "stdc++fs" "stdc++")
 
   ament_auto_package(INSTALL_TO_SHARE
     launch

launch/lightnet_trt.launch.xml

@@ -1,23 +1,26 @@
 <launch>
   <!-- image topic name to be subscribed -->
-  <arg name="input/image" default="image_raw" />
+  <arg name="input/image" default="/sensing/camera/camera0/image_rect_color"/>
+  <arg name="input/camera_info" default="/sensing/camera/camera0/camera_info"/>
+
   <!-- output topic name to be published -->
-  <arg name="output/objects" default="/perception/object_recognition/detection/rois0" />
+  <arg name="output/objects" default="/perception/object_recognition/detection/rois0"/>
 
   <!-- path to the lightNet-TRT cfg file to be loaded -->
-  <arg name="model_cfg" default="$(find-pkg-share lightnet_trt)/configs/lightNet-BDD100K-det-semaseg-1280x960.cfg" />
-  <arg name="model_weights" default="$(find-pkg-share lightnet_trt)/configs/lightNet-BDD100K-det-semaseg-1280x960.weights" />
-  <arg name="width" default="1280"/>
-  <arg name="height" default="960"/>
+  <!-- <arg name="model_cfg" default="$(find-pkg-share lightnet_trt)/configs/lightNet-BDD100K-det-semaseg-1280x960.cfg"/> -->
+  <!-- <arg name="model_weights" default="$(find-pkg-share lightnet_trt)/configs/lightNet-BDD100K-det-semaseg-1280x960.weights"/> -->
+  <!-- <arg name="width" default="1280"/> -->
+  <!-- <arg name="height" default="960"/> -->
 
-  <!-- <arg name="model_cfg" default="$(find-pkg-share lightnet_trt)/configs/yolov8x-BDD100K-640x640-T4-lane4cls.cfg"/> -->
-  <!-- <arg name="model_weights" default="$(find-pkg-share lightnet_trt)/configs/yolov8x-BDD100K-640x640-T4-lane4cls-scratch-pseudo-finetune_last.weights"/> -->
-  <!-- <arg name="width" default="640"/> -->
-  <!-- <arg name="height" default="640"/> -->
+  <arg name="model_cfg" default="$(find-pkg-share lightnet_trt)/configs/yolov8x-BDD100K-640x640-T4-lane4cls.cfg"/>
+  <arg name="model_weights" default="$(find-pkg-share lightnet_trt)/configs/yolov8x-BDD100K-640x640-T4-lane4cls-scratch-pseudo-finetune_last.weights"/>
+  <arg name="width" default="640"/>
+  <arg name="height" default="640"/>
 
   <group>
-    <node pkg="lightnet_trt" exec="lightnet_trt_node" name="lightnet_trt_node">
+    <node pkg="lightnet_trt" exec="lightnet_trt_node" name="lightnet_trt_node" output="screen">
       <remap from="~/in/image" to="$(var input/image)"/>
+      <remap from="~/in/camera_info" to="$(var input/camera_info)"/>
       <remap from="~/out/image" to="$(var output/objects)/debug/image"/>
       <param name="model_cfg" value="$(var model_cfg)"/>
       <param name="model_weights" value="$(var model_weights)"/>

modules/class_detector.cpp

@@ -50,6 +50,12 @@ void Detector::segment(const std::vector<cv::Mat> &mat_image, std::string filena
   _impl->_detector.segment(mat_image, filename);
 }
 
+void Detector::get_mask(std::vector<cv::Mat> &mask_results)
+{
+  _impl->_detector.get_mask(mask_results);
+}
+
+
 void Detector::draw_BBox(cv::Mat &img,Result result)
 {
   _impl->_detector.draw_BBox(img, result);

modules/class_detector.h

@@ -59,6 +59,7 @@ class API Detector
  void init(const Config &config);
  void save_image(cv::Mat &img, const std::string &dir, const std::string &name);
  void segment(const std::vector<cv::Mat> &mat_image, std::string filename);
+ void get_mask(std::vector<cv::Mat> &mask_results);
  void detect(const std::vector<cv::Mat> &mat_image, std::vector<BatchResult> &vec_batch_result, const bool cuda);
  void dump_profiling(void);
  void draw_BBox(cv::Mat &img, Result result); 

modules/class_yolo_detector.hpp

@@ -78,10 +78,8 @@ class YoloDectector
       cv::putText(img, stream.str(), cv::Point(result.rect.x, result.rect.y - 5), 0, 0.5, cv::Scalar(colormap[3*id+2], colormap[3*id+1], colormap[3*id+0]), 1);
     } else {
       cv::putText(img, stream.str(), cv::Point(result.rect.x, result.rect.y - 5), 0, 0.5, cv::Scalar(255, 0, 0), 1);
-    }
-
-
-  } 
+    } 
+  }
 
   void segment(const std::vector<cv::Mat> &vec_image, std::string filename)
   {
@@ -96,44 +94,50 @@ class YoloDectector
       int width = curImage.cols;
 
       for (uint32_t j = 0; j < mask.size(); j++) {
-	cv::Mat resized;
-
-	cv::resize(mask[j], resized, cv::Size(width, height), 0, 0, cv::INTER_NEAREST);	
-	cv::addWeighted(vec_image[i], 1.0, resized, 0.5, 0.0, vec_image[i]);
-
-	cv::namedWindow("mask" + std::to_string(j), cv::WINDOW_NORMAL);
-	cv::imshow("mask"+std::to_string(j), mask[j]);
-	if (flg_save) {
-	  fs::path p = save_path;
-	  p.append("segmentation");
-	  fs::create_directory(p);
-	  p.append(std::to_string(j));
-	  fs::create_directory(p);
-	  if (0) {
-	    //get filenames;
-	  } else {
-	    std::ostringstream sout;
-	    sout << std::setfill('0') << std::setw(6) << GLOBAL_COUNTER;
-	    //cv::resize(mask[j], resized, cv::Size(width, height), 0, 0, cv::INTER_NEAREST);
-	    if (filename == "") {
-	     std::string filename = "frame_" + sout.str() + ".png";
-	    }
-	    save_image(resized, p.string(), filename);
-	  }
-	}
-	auto depthmap = _p_net->get_depthmap(segmentation);
-	for (uint32_t j = 0; j < depthmap.size(); j++) {
-	  cv::imshow("depthmap"+std::to_string(j), depthmap[j]);
-	}
-      }
-      if (i > 0) {
-	GLOBAL_COUNTER++;
-      }
-    }
+				cv::Mat resized;
+				
+				cv::resize(mask[j], resized, cv::Size(width, height), 0, 0, cv::INTER_NEAREST);	
+				cv::addWeighted(vec_image[i], 1.0, resized, 0.5, 0.0, vec_image[i]);
+				
+				cv::namedWindow("mask" + std::to_string(j), cv::WINDOW_NORMAL);
+				cv::imshow("mask"+std::to_string(j), mask[j]);
+				if (flg_save) {
+					fs::path p = save_path;
+					p.append("segmentation");
+					fs::create_directory(p);
+					p.append(std::to_string(j));
+					fs::create_directory(p);
+					if (0) {
+						//get filenames;
+					} else {
+						std::ostringstream sout;
+						sout << std::setfill('0') << std::setw(6) << GLOBAL_COUNTER;
+						//cv::resize(mask[j], resized, cv::Size(width, height), 0, 0, cv::INTER_NEAREST);
+						if (filename == "") {
+							std::string filename = "frame_" + sout.str() + ".png";
+						}
+						save_image(resized, p.string(), filename);
+					}
+				}
+				auto depthmap = _p_net->get_depthmap(segmentation);
+				for (uint32_t j = 0; j < depthmap.size(); j++) {
+					cv::imshow("depthmap"+std::to_string(j), depthmap[j]);
+				}
+			}
+			if (i > 0) {
+				GLOBAL_COUNTER++;
+			}
+		}
     GLOBAL_COUNTER++;
-
   } 
 
+  void get_mask(std::vector<cv::Mat> &mask_results)
+  {
+		auto segmentation = _p_net->apply_argmax(0);
+		mask_results = _p_net->get_colorlbl(segmentation);
+    GLOBAL_COUNTER++;
+  }
+
   void dump_profiling(){
     _p_net->print_profiling();    
   }

package.xml

@@ -16,6 +16,8 @@
   <depend>cv_bridge</depend>
   <depend>image_transport</depend>
   <depend>rclcpp_components</depend>
+  <depend>tier4_autoware_utils</depend>
+  <depend>tf2_eigen</depend>
 
   <test_depend>ament_cmake_ros</test_depend>
   <test_depend>ament_lint_auto</test_depend>

samples/sample_detector.cpp

@@ -168,8 +168,8 @@ int main(int argc, char** argv)
 	  //	  std::cout <<"batch "<<i<< " id:" << r.id << " prob:" << r.prob << " rect:" << r.rect << std::endl;
 	  detector->draw_BBox(batch_img[i], r);
 	}
-	cv::namedWindow("image" + std::to_string(i), cv::WINDOW_NORMAL);
-	cv::imshow("image"+std::to_string(i), batch_img[i]);
+	// cv::namedWindow("image" + std::to_string(i), cv::WINDOW_NORMAL);
+	// cv::imshow("image"+std::to_string(i), batch_img[i]);
 	int k = cv::waitKey(0);
 	if (k == 32) {
 	  std::cout << "Save... " << name << std::endl;

src/lightnet_trt_core.cpp

@@ -22,20 +22,9 @@ LightNetTensorRT::LightNetTensorRT(const ::Config &config)
   detector_->init(config);
 }
 
-bool LightNetTensorRT::doInference(std::vector<cv::Mat> images)
+void LightNetTensorRT::doInference(const std::vector<cv::Mat> & images, std::vector<cv::Mat> & masks)
 {
   std::vector<BatchResult> batch_res;
-
   detector_->detect(images, batch_res, cuda);
-  detector_->segment(images, "");
-
-  for (int i = 0; i < images.size(); i++)
-  {
-    for (const auto &r : batch_res[i])
-    {
-      detector_->draw_BBox(images[i], r);
-    }
-  }
-
-  return true;
+  detector_->get_mask(masks);
 }

src/lightnet_trt_core.hpp

@@ -27,7 +27,7 @@ class LightNetTensorRT
 public:
   LightNetTensorRT(const ::Config &config);
 
-  bool doInference(std::vector<cv::Mat> images);
+  void doInference(const std::vector<cv::Mat> & images, std::vector<cv::Mat> & masks);
 
   std::unique_ptr<Detector> detector_;
   const bool cuda = false;

src/lightnet_trt_node.cpp

@@ -15,26 +15,31 @@
 #include "lightnet_trt_node.hpp"
 
 LightNetTensorRTNode::LightNetTensorRTNode(const rclcpp::NodeOptions &node_options)
-  : Node("tensorrt_lightnet", node_options)
+: Node("tensorrt_lightnet", node_options),
+  ipm_projector_(nullptr),
+  base_frame_("base_link")
 {
-  using std::placeholders::_1;
-  using std::chrono_literals::operator""ms;
+  transform_listener_ = std::make_shared<tier4_autoware_utils::TransformListener>(this);
 
   // Initialize a detector
-  ::Config config = loadConfig();
+  config_ = loadConfig();
 
   RCLCPP_INFO(this->get_logger(), "Start loading YOLO");
-  RCLCPP_INFO(this->get_logger(), "Model Config: %s", config.file_model_cfg.c_str());
-  RCLCPP_INFO(this->get_logger(), "Model Weights: %s", config.file_model_weights.c_str());
-  RCLCPP_INFO(this->get_logger(), "Input size: (%d, %d)", config.width, config.height);
-  lightnet_trt_ = std::make_unique<LightNetTensorRT>(config);
+  RCLCPP_INFO(this->get_logger(), "Model Config: %s", config_.file_model_cfg.c_str());
+  RCLCPP_INFO(this->get_logger(), "Model Weights: %s", config_.file_model_weights.c_str());
+  RCLCPP_INFO(this->get_logger(), "Input size: (%d, %d)", config_.width, config_.height);
+  lightnet_trt_ = std::make_unique<LightNetTensorRT>(config_);
   RCLCPP_INFO(this->get_logger(), "Finished loading YOLO");
 
   image_sub_ = image_transport::create_subscription(
-    this, "~/in/image", std::bind(&LightNetTensorRTNode::onImage, this, _1), "raw",
+    this, "~/in/image", [this](const sensor_msgs::msg::Image::ConstSharedPtr msg) { onImage(msg); }, "raw",
     rmw_qos_profile_sensor_data);
+  camera_info_sub_ = this->create_subscription<sensor_msgs::msg::CameraInfo>(
+    "~/in/camera_info", rclcpp::SensorDataQoS(),
+    [this](const sensor_msgs::msg::CameraInfo::ConstSharedPtr msg) { onCameraInfo(msg); });
 
   image_pub_ = image_transport::create_publisher(this, "~/out/image");
+  image_projected_pub_ = image_transport::create_publisher(this, "~/out/image_projected");
 }
 
 ::Config LightNetTensorRTNode::loadConfig()
@@ -59,19 +64,56 @@ ::Config LightNetTensorRTNode::loadConfig()
   return config;
 }
 
+void LightNetTensorRTNode::onCameraInfo(const sensor_msgs::msg::CameraInfo::ConstSharedPtr msg)
+{
+  if (ipm_projector_ != nullptr) return;
+  const Eigen::Matrix3f intrinsic = get_intrinsic_matrix(*msg);
+  const Eigen::Matrix3f intrinsic_resized = resize_intrinsic_matrix(intrinsic, config_.width, config_.height);
+
+  const std::string camera_frame = msg->header.frame_id;
+  geometry_msgs::msg::TransformStamped::ConstSharedPtr tf_base2cam_ptr =
+    transform_listener_->getLatestTransform(base_frame_, camera_frame);
+  if (!tf_base2cam_ptr) return;
+
+  const Eigen::Matrix4f extrinsic = tf2::transformToEigen(tf_base2cam_ptr->transform).matrix().cast<float>();
+  ipm_projector_ = std::make_unique<IPM>(intrinsic_resized, extrinsic);
+  return;
+}
+
 void LightNetTensorRTNode::onImage(const sensor_msgs::msg::Image::ConstSharedPtr msg)
 {
   RCLCPP_INFO(this->get_logger(), "Image received");
   cv_bridge::CvImagePtr in_image_ptr;
   in_image_ptr = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::BGR8);
 
-  if (!lightnet_trt_->doInference({in_image_ptr->image}))
+  std::vector<cv::Mat> masks;
+  lightnet_trt_->doInference({in_image_ptr->image}, masks);
+  RCLCPP_INFO(this->get_logger(), "Inference succeeded");
+
+  // Publish result
+  cv_bridge::CvImage out_image;
+  out_image.header = msg->header;
+  out_image.encoding = sensor_msgs::image_encodings::BGR8;
+  out_image.image = masks[0];
+  image_pub_.publish(out_image.toImageMsg());
+
+  // Post processing
+  if (!ipm_projector_)
   {
-    RCLCPP_WARN(this->get_logger(), "Inference failed.");
+    RCLCPP_WARN(this->get_logger(), "IPM projector is not initialized.");
     return;
   }
-  image_pub_.publish(in_image_ptr->toImageMsg());
-  RCLCPP_INFO(this->get_logger(), "Inference succeeded");
+  RCLCPP_WARN(this->get_logger(), "IPM projector running.");
+
+  cv::Mat projected_mask;
+  ipm_projector_->run(masks[0], projected_mask);
+
+  // Publish result
+  cv_bridge::CvImage out_projected_image;
+  out_projected_image.header = msg->header;
+  out_projected_image.encoding = sensor_msgs::image_encodings::BGR8;
+  out_projected_image.image = projected_mask;
+  image_projected_pub_.publish(out_projected_image.toImageMsg());
 }
 
 #include "rclcpp_components/register_node_macro.hpp"

src/lightnet_trt_node.hpp

@@ -15,13 +15,18 @@
 #ifndef LIGHTNET_TRT__LIGHTNET_TRT_NODE_HPP_
 #define LIGHTNET_TRT__LIGHTNET_TRT_NODE_HPP_
 
-#include <image_transport/image_transport.hpp>
-#include <rclcpp/rclcpp.hpp>
 #include "lightnet_trt_core.hpp"
+#include "utils.hpp"
 
-#include <sensor_msgs/msg/image.hpp>
+#include <rclcpp/rclcpp.hpp>
 
 #include <cv_bridge/cv_bridge.h>
+#include <image_transport/image_transport.hpp>
+#include <sensor_msgs/msg/image.hpp>
+#include <tf2_eigen/tf2_eigen.hpp>
+#include <tier4_autoware_utils/ros/transform_listener.hpp>
+
+#include <optional>
 
 class LightNetTensorRTNode : public rclcpp::Node
 {
@@ -30,10 +35,23 @@ class LightNetTensorRTNode : public rclcpp::Node
 
 private:
   void onImage(const sensor_msgs::msg::Image::ConstSharedPtr msg);
+  void onCameraInfo(const sensor_msgs::msg::CameraInfo::ConstSharedPtr msg);
   ::Config loadConfig();
 
-  image_transport::Publisher image_pub_;
+  // ROS related variables
   image_transport::Subscriber image_sub_;
+  rclcpp::Subscription<sensor_msgs::msg::CameraInfo>::SharedPtr camera_info_sub_;
+  image_transport::Publisher image_pub_;
+  image_transport::Publisher image_projected_pub_;
+  std::shared_ptr<tier4_autoware_utils::TransformListener> transform_listener_;
+
+  // Tools
   std::unique_ptr<LightNetTensorRT> lightnet_trt_;
+  std::unique_ptr<IPM> ipm_projector_;
+
+  // Misc
+  bool received_camera_info_;
+  const std::string base_frame_;
+  ::Config config_;
 };
 #endif  // LIGHTNET_TRT__LIGHTNET_TRT_NODE_HPP_