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

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
 )
 
@@ -166,6 +167,12 @@ rclcpp_components_register_node(${PROJECT_NAME}_node
 if(BUILD_TESTING)
   find_package(ament_lint_auto REQUIRED)
   ament_lint_auto_find_test_dependencies()
+
+  if(TRT_AVAIL AND CUDA_AVAIL AND CUDNN_AVAIL)
+    ament_auto_add_gtest(test_utils
+      test/test_utils.cpp
+    )
+  endif()
 endif()
 
 ament_auto_package(INSTALL_TO_SHARE

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_;

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_

perception/autoware_tensorrt_yolox/package.xml

@@ -32,6 +32,7 @@
 
   <exec_depend>autoware_image_transport_decompressor</exec_depend>
 
+  <test_depend>ament_cmake_gtest</test_depend>
   <test_depend>ament_lint_auto</test_depend>
   <test_depend>autoware_lint_common</test_depend>
 

perception/autoware_tensorrt_yolox/src/tensorrt_yolox_node.cpp

@@ -1,4 +1,4 @@
 // Copyright 2022 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.
@@ -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>
@@ -176,13 +177,20 @@
     if (is_roi_overlap_segment_ && trt_yolox_->getMultitaskNum() > 0) {
       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());

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

perception/autoware_tensorrt_yolox/test/test_utils.cpp

@@ -0,0 +1,80 @@
+// 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"
+
+#include <opencv2/opencv.hpp>
+
+#include <gtest/gtest.h>
+
+using autoware::tensorrt_yolox::runLengthDecoder;
+using autoware::tensorrt_yolox::runLengthEncoder;
+
+// Test case 1: Test if the decoded image is the same as the original image
+TEST(UtilsTest, runLengthEncoderDecoderTest)
+{
+  int height = 10;
+  int width = 20;
+  uint8_t number_cls = 16;
+  // Create an image as below
+  // 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
+  // 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
+  // 2 2 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
+  // 3 3 3 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
+  // 4 4 4 4 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
+  // 5 5 5 5 5 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
+  // 6 6 6 6 6 6 0 0 0 0 0 0 0 0 0 0 0 0 0 0
+  // 7 7 7 7 7 7 7 0 0 0 0 0 0 0 0 0 0 0 0 0
+  // 8 8 8 8 8 8 8 8 0 0 0 0 0 0 0 0 0 0 0 0
+  // 9 9 9 9 9 9 9 9 9 0 0 0 0 0 0 0 0 0 0 0
+  cv::Mat image = cv::Mat::zeros(10, 20, CV_8UC1);
+  for (int i = 0; i < height; ++i) {
+    for (int j = 0; j < width; ++j) {
+      if (j < i) {
+        image.at<uint8_t>(i, j) = i % number_cls;
+      } else {
+        image.at<uint8_t>(i, j) = 0;
+      }
+    }
+  }
+  // Compress the image
+  std::vector<std::pair<uint8_t, int>> compressed_data = runLengthEncoder(image);
+  int step = sizeof(uint8_t) + sizeof(int);
+  std::vector<uint8_t> data(compressed_data.size() * step);
+  for (size_t i = 0; i < compressed_data.size(); ++i) {
+    std::memcpy(&data[i * step], &compressed_data.at(i).first, sizeof(uint8_t));
+    std::memcpy(&data[i * step + sizeof(uint8_t)], &compressed_data.at(i).second, sizeof(int));
+  }
+  // Decompress the image
+  cv::Mat decoded_image = runLengthDecoder(data, height, width);
+  // Compare the original image and the decoded image
+  ASSERT_EQ(image.rows, decoded_image.rows);
+  ASSERT_EQ(image.cols, decoded_image.cols);
+  bool image_eq = true;
+  for (int i = 0; i < image.rows; ++i) {
+    for (int j = 0; j < image.cols; ++j) {
+      if (image.at<uint8_t>(i, j) != decoded_image.at<uint8_t>(i, j)) {
+        image_eq = false;
+        break;
+      }
+    }
+  }
+  EXPECT_EQ(image_eq, true);
+}
+
+int main(int argc, char ** argv)
+{
+  testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}