fix(tensorrt_yolox): add run length encoding for sematic segmentation…

… mask (autowarefoundation#7905) * fix: add rle compress Signed-off-by: badai-nguyen <[email protected]> * fix: rle compress Signed-off-by: badai-nguyen <[email protected]> * fix: move rle into utils Signed-off-by: badai-nguyen <[email protected]> * chore: pre-commit Signed-off-by: badai-nguyen <[email protected]> * Update perception/autoware_tensorrt_yolox/src/utils.cpp Co-authored-by: Yukihiro Saito <[email protected]> * fix: remove unused variable Signed-off-by: badai-nguyen <[email protected]> * Update perception/autoware_tensorrt_yolox/src/utils.cpp Co-authored-by: Manato Hirabayashi <[email protected]> * style(pre-commit): autofix * feat: add unit test for utils Signed-off-by: badai-nguyen <[email protected]> * style(pre-commit): autofix * fix: unit test Signed-off-by: badai-nguyen <[email protected]> * chore: change to explicit index Signed-off-by: badai-nguyen <[email protected]> * style(pre-commit): autofix * fix: cuda cmake Signed-off-by: badai-nguyen <[email protected]> * fix: separate unit test into different PR Signed-off-by: badai-nguyen <[email protected]> --------- Signed-off-by: badai-nguyen <[email protected]> Co-authored-by: Yukihiro Saito <[email protected]> Co-authored-by: Manato Hirabayashi <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
esteve · Aug 13, 2024 · 160493a · 160493a
1 parent 775576c
commit 160493a
Show file tree

Hide file tree

Showing 5 changed files with 105 additions and 1 deletion.
diff --git a/perception/autoware_tensorrt_yolox/CMakeLists.txt b/perception/autoware_tensorrt_yolox/CMakeLists.txt
@@ -143,6 +143,7 @@ rclcpp_components_register_node(yolox_single_image_inference_node
 )
 
 ament_auto_add_library(${PROJECT_NAME}_node SHARED
+  src/utils.cpp
   src/tensorrt_yolox_node.cpp
 )
 

diff --git a/perception/autoware_tensorrt_yolox/include/autoware/tensorrt_yolox/tensorrt_yolox_node.hpp b/perception/autoware_tensorrt_yolox/include/autoware/tensorrt_yolox/tensorrt_yolox_node.hpp
@@ -40,6 +40,7 @@
 #include <map>
 #include <memory>
 #include <string>
+#include <utility>
 #include <vector>
 
 namespace autoware::tensorrt_yolox
@@ -81,6 +82,7 @@ class TrtYoloXNode : public rclcpp::Node
   int mapRoiLabel2SegLabel(const int32_t roi_label_index);
   image_transport::Publisher image_pub_;
   image_transport::Publisher mask_pub_;
+
   image_transport::Publisher color_mask_pub_;
 
   rclcpp::Publisher<tier4_perception_msgs::msg::DetectedObjectsWithFeature>::SharedPtr objects_pub_;

diff --git a/perception/autoware_tensorrt_yolox/include/autoware/tensorrt_yolox/utils.hpp b/perception/autoware_tensorrt_yolox/include/autoware/tensorrt_yolox/utils.hpp
@@ -0,0 +1,28 @@
+// Copyright 2024 TIER IV, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef AUTOWARE__TENSORRT_YOLOX__UTILS_HPP_
+#define AUTOWARE__TENSORRT_YOLOX__UTILS_HPP_
+#include <opencv2/opencv.hpp>
+
+#include <utility>
+#include <vector>
+
+namespace autoware::tensorrt_yolox
+{
+std::vector<std::pair<uint8_t, int>> runLengthEncoder(const cv::Mat & mask);
+cv::Mat runLengthDecoder(const std::vector<uint8_t> & rle_data, const int rows, const int cols);
+}  // namespace autoware::tensorrt_yolox
+
+#endif  // AUTOWARE__TENSORRT_YOLOX__UTILS_HPP_
diff --git a/perception/autoware_tensorrt_yolox/src/tensorrt_yolox_node.cpp b/perception/autoware_tensorrt_yolox/src/tensorrt_yolox_node.cpp
@@ -14,6 +14,7 @@
 
 #include "autoware/tensorrt_yolox/tensorrt_yolox_node.hpp"
 
+#include "autoware/tensorrt_yolox/utils.hpp"
 #include "object_recognition_utils/object_classification.hpp"
 
 #include <autoware_perception_msgs/msg/object_classification.hpp>
@@ -177,12 +178,19 @@ void TrtYoloXNode::onImage(const sensor_msgs::msg::Image::ConstSharedPtr msg)
       overlapSegmentByRoi(yolox_object, mask, width, height);
     }
   }
-  // TODO(badai-nguyen): consider to change to 4bits data transfer
   if (trt_yolox_->getMultitaskNum() > 0) {
     sensor_msgs::msg::Image::SharedPtr out_mask_msg =
       cv_bridge::CvImage(std_msgs::msg::Header(), sensor_msgs::image_encodings::MONO8, mask)
         .toImageMsg();
     out_mask_msg->header = msg->header;
+
+    std::vector<std::pair<uint8_t, int>> compressed_data = runLengthEncoder(mask);
+    int step = sizeof(uint8_t) + sizeof(int);
+    out_mask_msg->data.resize(static_cast<int>(compressed_data.size()) * step);
+    for (size_t i = 0; i < compressed_data.size(); ++i) {
+      std::memcpy(&out_mask_msg->data[i * step], &compressed_data.at(i).first, sizeof(uint8_t));
+      std::memcpy(&out_mask_msg->data[i * step + 1], &compressed_data.at(i).second, sizeof(int));
+    }
     mask_pub_.publish(out_mask_msg);
   }
   image_pub_.publish(in_image_ptr->toImageMsg());

diff --git a/perception/autoware_tensorrt_yolox/src/utils.cpp b/perception/autoware_tensorrt_yolox/src/utils.cpp
@@ -0,0 +1,65 @@
+// Copyright 2024 TIER IV, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "autoware/tensorrt_yolox/utils.hpp"
+
+namespace autoware::tensorrt_yolox
+{
+
+std::vector<std::pair<uint8_t, int>> runLengthEncoder(const cv::Mat & image)
+{
+  std::vector<std::pair<uint8_t, int>> compressed_data;
+  const int rows = image.rows;
+  const int cols = image.cols;
+  compressed_data.emplace_back(image.at<uint8_t>(0, 0), 0);
+  for (int i = 0; i < rows; ++i) {
+    for (int j = 0; j < cols; ++j) {
+      uint8_t current_value = image.at<uint8_t>(i, j);
+      if (compressed_data.back().first == current_value) {
+        ++compressed_data.back().second;
+      } else {
+        compressed_data.emplace_back(current_value, 1);
+      }
+    }
+  }
+  return compressed_data;
+}
+
+cv::Mat runLengthDecoder(const std::vector<uint8_t> & rle_data, const int rows, const int cols)
+{
+  cv::Mat mask(rows, cols, CV_8UC1, cv::Scalar(0));
+  int idx = 0;
+  int step = sizeof(uint8_t) + sizeof(int);
+  for (size_t i = 0; i < rle_data.size(); i += step) {
+    uint8_t value;
+    int length;
+    std::memcpy(&value, &rle_data[i], sizeof(uint8_t));
+    std::memcpy(
+      &length, &rle_data[i + 1],
+      sizeof(
+        int));  // under the condition that we know rle_data[i] only consume 1 element of the vector
+    for (int j = 0; j < length; ++j) {
+      int row_idx = static_cast<int>(idx / cols);
+      int col_idx = static_cast<int>(idx % cols);
+      mask.at<uint8_t>(row_idx, col_idx) = value;
+      idx++;
+      if (idx > rows * cols) {
+        break;
+      }
+    }
+  }
+  return mask;
+}
+
+}  // namespace autoware::tensorrt_yolox