chore: refactor

Signed-off-by: badai-nguyen <[email protected]>

perception/tensorrt_yolox/CMakeLists.txt

@@ -114,4 +114,5 @@ endif()
 
 ament_auto_package(INSTALL_TO_SHARE
   launch
+  config
 )

perception/tensorrt_yolox/config/yolox_s_plus_opt.param.yaml

@@ -0,0 +1,20 @@
+/**:
+  ros__parameters:
+    # refine segmentation mask by overlay roi class
+    # disable when sematic segmentation accuracy is good enough
+    is_roi_overlap_segment: true
+
+    # minimum existence_probability of detected roi considered to replace segmentation
+    overlap_roi_score_threshold: 0.3
+
+    # publish color mask for result visualization
+    is_publish_color_mask: false
+
+    roi_overlay_segment_label:
+      UNKNOWN : true
+      CAR : true
+      TRUCK : true
+      BUS : true
+      MOTORCYCLE : true
+      BICYCLE : true
+      PEDESTRIAN : true

perception/tensorrt_yolox/include/tensorrt_yolox/tensorrt_yolox.hpp

@@ -78,7 +78,6 @@ class TrtYoloX
    * @param[in] build_config configuration including precision, calibration method, DLA, remaining
    * fp16 for first layer,  remaining fp16 for last layer and profiler for builder
    * @param[in] use_gpu_preprocess whether use cuda gpu for preprocessing
-   * @param[in] publish_color_mask whether publish color_mask for debugging and visualization
    * @param[in] calibration_image_list_file path for calibration files (only require for
    * quantization)
    * @param[in] norm_factor scaling factor for preprocess
@@ -91,9 +90,8 @@ class TrtYoloX
     const std::string & color_map_path, const int num_class = 8, const float score_threshold = 0.3,
     const float nms_threshold = 0.7,
     const tensorrt_common::BuildConfig build_config = tensorrt_common::BuildConfig(),
-    const bool use_gpu_preprocess = false, const bool publish_color_mask = false,
-    std::string calibration_image_list_file = std::string(), const double norm_factor = 1.0,
-    [[maybe_unused]] const std::string & cache_dir = "",
+    const bool use_gpu_preprocess = false, std::string calibration_image_list_file = std::string(),
+    const double norm_factor = 1.0, [[maybe_unused]] const std::string & cache_dir = "",
     const tensorrt_common::BatchConfig & batch_config = {1, 1, 1},
     const size_t max_workspace_size = (1 << 30));
   /**
@@ -153,7 +151,7 @@ class TrtYoloX
    * @param[in] index multitask index
    * @param[in] colormap colormap for masks
    */
-    void getColorizedMask(
+  void getColorizedMask(
     const std::vector<tensorrt_yolox::Colormap> & colormap, const cv::Mat & mask,
     cv::Mat & colorized_mask);
   inline std::vector<Colormap> getColorMap() { return sematic_color_map_; }
@@ -313,10 +311,6 @@ class TrtYoloX
   // device buffer for argmax postprocessing  on GPU
   CudaUniquePtr<unsigned char[]> argmax_buf_d_;
   std::vector<tensorrt_yolox::Colormap> sematic_color_map_;
-  // flag whether overlay segmentation by roi
-  bool roi_overlap_segment_;
-  // flag where publish color mask for debugging and visualization
-  bool publish_color_mask_;
 };
 
 }  // namespace tensorrt_yolox

perception/tensorrt_yolox/include/tensorrt_yolox/tensorrt_yolox_node.hpp

@@ -15,6 +15,9 @@
 #ifndef TENSORRT_YOLOX__TENSORRT_YOLOX_NODE_HPP_
 #define TENSORRT_YOLOX__TENSORRT_YOLOX_NODE_HPP_
 
+#include "object_recognition_utils/object_recognition_utils.hpp"
+#include "utils/utils.hpp"
+
 #include <image_transport/image_transport.hpp>
 #include <opencv2/opencv.hpp>
 #include <rclcpp/rclcpp.hpp>
@@ -68,6 +71,18 @@ class TrtYoloXNode : public rclcpp::Node
   bool is_roi_overlap_segment_;
   bool is_publish_color_mask_;
   float overlap_roi_score_threshold_;
+  // TODO(badai-nguyen): change to function
+  std::map<std::string, int> remap_roi_to_semantic_ = {
+    {"UNKNOWN", 19},     // other
+    {"ANIMAL", 19},      // other
+    {"PEDESTRIAN", 11},  // person
+    {"CAR", 13},         // car
+    {"TRUCK", 14},       // truck
+    {"BUS", 15},         // bus
+    {"BICYCLE", 18},     // bicycle
+    {"MOTORBIKE", 17},   // motorcycle
+  };
+  utils::FilterTargetLabel roi_overlay_segment_labels_;
 };
 
 }  // namespace tensorrt_yolox

perception/tensorrt_yolox/launch/yolox_s_plus_opt.launch.xml

@@ -32,9 +32,7 @@
   <arg name="calibration_image_list_path" default="" description="Path to a file which contains path to images. Those images will be used for int8 quantization."/>
   <arg name="use_decompress" default="true" description="use image decompress"/>
   <arg name="build_only" default="false" description="exit after trt engine is built"/>
-  <arg name="is_roi_overlap_segment" default="true" description="refine segmentation mask by overlay roi class, disable when sematic segmentation accuracy is good enough"/>
-  <arg name="is_publish_color_mask" default="false" description="publish color mask for result visualization"/>
-  <arg name="overlap_roi_score_threshold" default="0.3" description="the roi object existance probability threshold that consider to replace segmentation"/>
+  <arg name="yolox_s_plus_opt_param_path" default="$(find-pkg-share tensorrt_yolox)/config/yolox_s_plus_opt.param.yaml"/>
   <node pkg="image_transport_decompressor" exec="image_transport_decompressor_node" name="image_transport_decompressor_node" if="$(var use_decompress)">
     <remap from="~/input/compressed_image" to="$(var input/image)/compressed"/>
     <remap from="~/output/raw_image" to="$(var input/image)"/>
@@ -59,8 +57,6 @@
     <param name="calibration_image_list_path" value="$(var calibration_image_list_path)"/>
     <param name="build_only" value="$(var build_only)"/>
     <param name="color_map_path" value="$(var model_path)/semseg_color_map.csv"/>
-    <param name="is_roi_overlap_segment" value="$(var is_roi_overlap_segment)"/>
-    <param name="is_publish_color_mask" value="$(var is_publish_color_mask)"/>
-    <param name="overlap_roi_score_threshold" value="$(var overlap_roi_score_threshold)"/>
+    <param from="$(var yolox_s_plus_opt_param_path)"/>
   </node>
 </launch>

perception/tensorrt_yolox/package.xml

@@ -22,6 +22,7 @@
   <depend>autoware_auto_perception_msgs</depend>
   <depend>cuda_utils</depend>
   <depend>cv_bridge</depend>
+  <depend>detected_object_validation</depend>
   <depend>image_transport</depend>
   <depend>libopencv-dev</depend>
   <depend>object_recognition_utils</depend>

perception/tensorrt_yolox/src/tensorrt_yolox.cpp

@@ -157,7 +157,7 @@ namespace tensorrt_yolox
 TrtYoloX::TrtYoloX(
   const std::string & model_path, const std::string & precision, const std::string & color_map_path,
   const int num_class, const float score_threshold, const float nms_threshold,
-  tensorrt_common::BuildConfig build_config, const bool use_gpu_preprocess, bool publish_color_mask,
+  tensorrt_common::BuildConfig build_config, const bool use_gpu_preprocess,
   std::string calibration_image_list_path, const double norm_factor,
   [[maybe_unused]] const std::string & cache_dir, const tensorrt_common::BatchConfig & batch_config,
   const size_t max_workspace_size)
@@ -168,7 +168,6 @@ TrtYoloX::TrtYoloX(
   batch_size_ = batch_config[2];
   multitask_ = 0;
   sematic_color_map_ = get_seg_colormap(color_map_path);
-  publish_color_mask_ = publish_color_mask;
   if (precision == "int8") {
     if (build_config.clip_value <= 0.0) {
       if (calibration_image_list_path.empty()) {
@@ -1281,7 +1280,10 @@ void TrtYoloX::getColorizedMask(
 {
   int width = mask.cols;
   int height = mask.rows;
-  // TODO: check size of mask and cmask
+  if ((cmask.cols != mask.cols) || (cmask.rows != mask.rows)) {
+    throw std::runtime_error("input and output image have difference size.");
+    return;
+  }
   for (int y = 0; y < height; y++) {
     for (int x = 0; x < width; x++) {
       unsigned char id = mask.at<unsigned char>(y, x);

perception/tensorrt_yolox/src/tensorrt_yolox_node.cpp

@@ -98,6 +98,17 @@ TrtYoloXNode::TrtYoloXNode(const rclcpp::NodeOptions & node_options)
   is_roi_overlap_segment_ = declare_parameter<bool>("is_roi_overlap_segment");
   is_publish_color_mask_ = declare_parameter<bool>("is_publish_color_mask");
   overlap_roi_score_threshold_ = declare_parameter<float>("overlap_roi_score_threshold");
+  roi_overlay_segment_labels_.UNKNOWN =
+    declare_parameter<bool>("roi_overlay_segment_label.UNKNOWN");
+  roi_overlay_segment_labels_.CAR = declare_parameter<bool>("roi_overlay_segment_label.CAR");
+  roi_overlay_segment_labels_.TRUCK = declare_parameter<bool>("roi_overlay_segment_label.TRUCK");
+  roi_overlay_segment_labels_.BUS = declare_parameter<bool>("roi_overlay_segment_label.BUS");
+  roi_overlay_segment_labels_.MOTORCYCLE =
+    declare_parameter<bool>("roi_overlay_segment_label.MOTORCYCLE");
+  roi_overlay_segment_labels_.BICYCLE =
+    declare_parameter<bool>("roi_overlay_segment_label.BICYCLE");
+  roi_overlay_segment_labels_.PEDESTRIAN =
+    declare_parameter<bool>("roi_overlay_segment_label.PEDESTRIAN");
   replaceLabelMap();
 
   tensorrt_common::BuildConfig build_config(
@@ -106,7 +117,7 @@ TrtYoloXNode::TrtYoloXNode(const rclcpp::NodeOptions & node_options)
 
   trt_yolox_ = std::make_unique<tensorrt_yolox::TrtYoloX>(
     model_path, precision, color_map_path, label_map_.size(), score_threshold, nms_threshold,
-    build_config, preprocess_on_gpu, is_publish_color_mask_, calibration_image_list_path);
+    build_config, preprocess_on_gpu, calibration_image_list_path);
 
   timer_ =
     rclcpp::create_timer(this, get_clock(), 100ms, std::bind(&TrtYoloXNode::onConnect, this));
@@ -163,6 +174,7 @@ void TrtYoloXNode::onImage(const sensor_msgs::msg::Image::ConstSharedPtr msg)
     return;
   }
   auto & mask = masks.at(0);
+  // TODO(badai-nguyen): change to postprocess on gpu option
   cv::resize(
     mask, mask, cv::Size(in_image_ptr->image.cols, in_image_ptr->image.rows), 0, 0,
     cv::INTER_NEAREST);
@@ -192,6 +204,7 @@ void TrtYoloXNode::onImage(const sensor_msgs::msg::Image::ConstSharedPtr msg)
       overlapSegmentByRoi(yolox_object, mask);
     }
   }
+  // TODO(badai-nguyen): consider to change to 4bits data transfer
   sensor_msgs::msg::Image::SharedPtr out_mask_msg =
     cv_bridge::CvImage(std_msgs::msg::Header(), sensor_msgs::image_encodings::MONO8, mask)
       .toImageMsg();
@@ -250,23 +263,10 @@ void TrtYoloXNode::replaceLabelMap()
 
 int TrtYoloXNode::mapRoiLabel2SegLabel(const int32_t roi_label_index)
 {
-  switch (roi_label_index) {
-    case 0:
-      return 5;
-    case 1:
-      return 13;
-    case 2:
-      return 14;
-    case 3:
-      return 15;
-    case 4:
-      return 18;
-    case 5:
-      return 17;
-    case 6:
-      return 11;
-    default:
-      return -1;
+  if (roi_overlay_segment_labels_.isTarget(static_cast<uint8_t>(roi_label_index))) {
+    std::string label = label_map_[roi_label_index];
+
+    return remap_roi_to_semantic_[label];
   }
   return -1;
 }