fix: roi overlay segment

Signed-off-by: badai-nguyen <dai.nguyen@tier4.jp>

Author: badai-nguyen

perception/tensorrt_yolox/include/tensorrt_yolox/tensorrt_yolox.hpp

@@ -153,7 +153,10 @@ class TrtYoloX
    * @param[in] index multitask index
    * @param[in] colormap colormap for masks
    */
-  cv::Mat getColorizedMask(int index, std::vector<Colormap> & colormap);
+    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_; }
 
 private:
   /**

perception/tensorrt_yolox/src/tensorrt_yolox.cpp

@@ -918,7 +918,6 @@ bool TrtYoloX::feedforwardAndDecode(
   for (size_t i = 0; i < batch_size; ++i) {
     auto image_size = images[i].size();
     auto & out_mask = out_masks[i];
-    auto & color_mask = color_masks[i];
     float * batch_prob = out_prob_h_.get() + (i * out_elem_num_per_batch_);
     ObjectArray object_array;
     decodeOutputs(batch_prob, object_array, scales_[i], image_size);
@@ -957,13 +956,7 @@ bool TrtYoloX::feedforwardAndDecode(
       continue;
     }
     // Assume semantic segmentation is first task
-    // This should remove when the segmentation accuracy is high
     out_mask = segmentation_masks_.at(0);
-
-    // publish color mask for visualization
-    if (publish_color_mask_) {
-      color_mask = getColorizedMask(0, sematic_color_map_);
-    }
   }
   return true;
 }
@@ -1283,13 +1276,12 @@ int TrtYoloX::getMultitaskNum(void)
   return multitask_;
 }
 
-cv::Mat TrtYoloX::getColorizedMask(int index, std::vector<Colormap> & colormap)
+void TrtYoloX::getColorizedMask(
+  const std::vector<tensorrt_yolox::Colormap> & colormap, const cv::Mat & mask, cv::Mat & cmask)
 {
-  cv::Mat mask;
-  mask = segmentation_masks_[index];
   int width = mask.cols;
   int height = mask.rows;
-  cv::Mat cmask = cv::Mat::zeros(height, width, CV_8UC3);
+  // TODO: check size of mask and cmask
   for (int y = 0; y < height; y++) {
     for (int x = 0; x < width; x++) {
       unsigned char id = mask.at<unsigned char>(y, x);
@@ -1298,7 +1290,6 @@ cv::Mat TrtYoloX::getColorizedMask(int index, std::vector<Colormap> & colormap)
       cmask.at<cv::Vec3b>(y, x)[2] = colormap[id].color[0];
     }
   }
-  return cmask;
 }
 
 }  // namespace tensorrt_yolox

perception/tensorrt_yolox/src/tensorrt_yolox_node.cpp

@@ -187,6 +187,7 @@ void TrtYoloXNode::onImage(const sensor_msgs::msg::Image::ConstSharedPtr msg)
       in_image_ptr->image, cv::Point(left, top), cv::Point(right, bottom), cv::Scalar(0, 0, 255), 3,
       8, 0);
     // Refine mask: replacing segmentation mask by roi class
+    // This should remove when the segmentation accuracy is high
     if (is_roi_overlap_segment_) {
       overlapSegmentByRoi(yolox_object, mask);
     }
@@ -202,10 +203,9 @@ void TrtYoloXNode::onImage(const sensor_msgs::msg::Image::ConstSharedPtr msg)
   objects_pub_->publish(out_objects);
 
   if (is_publish_color_mask_) {
-    auto & color_mask = color_masks.at(0);
-    cv::resize(
-      color_mask, color_mask, cv::Size(in_image_ptr->image.cols, in_image_ptr->image.rows), 0, 0,
-      cv::INTER_NEAREST);
+    cv::Mat color_mask =
+      cv::Mat::zeros(in_image_ptr->image.rows, in_image_ptr->image.cols, CV_8UC3);
+    trt_yolox_->getColorizedMask(trt_yolox_->getColorMap(), mask, color_mask);
     sensor_msgs::msg::Image::SharedPtr output_color_mask_msg =
       cv_bridge::CvImage(std_msgs::msg::Header(), sensor_msgs::image_encodings::BGR8, color_mask)
         .toImageMsg();
@@ -250,31 +250,37 @@ void TrtYoloXNode::replaceLabelMap()
 
 int TrtYoloXNode::mapRoiLabel2SegLabel(const int32_t roi_label_index)
 {
-  auto & roi_label = label_map_[roi_label_index];
-  if (roi_label == "CAR" || roi_label == "BUS" || roi_label == "TRUCK") {
-    return static_cast<int>(roi_label_index + 11);
-  }
-  if (roi_label == "PEDESTRIAN") {
-    return 11;  // person index in segment_color_map
-  }
-  if (roi_label == "MOTORCYCLE") {
-    return 17;  // motocycle index in segment_color_map
-  }
-  if (roi_label == "BICYCLE") {
-    return 18;  // bicycle index in segment_color_map
+  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;
   }
   return -1;
 }
 
 void TrtYoloXNode::overlapSegmentByRoi(const tensorrt_yolox::Object & roi_object, cv::Mat & mask)
 {
   if (roi_object.score < overlap_roi_score_threshold_) return;
-  cv::Mat submat = mask.colRange(roi_object.x_offset, roi_object.width)
-                     .rowRange(roi_object.y_offset, roi_object.height);
   int seg_class_index = mapRoiLabel2SegLabel(roi_object.type);
   if (seg_class_index < 0) return;
-  cv::Mat replace_roi(cv::Size(), mask.type(), seg_class_index);
-  replace_roi.copyTo(submat);
+  cv::Mat replace_roi(
+    cv::Size(roi_object.width, roi_object.height), mask.type(),
+    static_cast<uint8_t>(seg_class_index));
+  replace_roi.copyTo(mask.colRange(roi_object.x_offset, roi_object.x_offset + roi_object.width)
+                       .rowRange(roi_object.y_offset, roi_object.y_offset + roi_object.height));
 }
 
 }  // namespace tensorrt_yolox