From dcd56fe9e20e641bc9b392de8a82a18c9b5f1b3e Mon Sep 17 00:00:00 2001 From: badai nguyen <94814556+badai-nguyen@users.noreply.github.com> Date: Thu, 8 Aug 2024 22:05:44 +0900 Subject: [PATCH] fix(tensorrt_yolox): add run length encoding for sematic segmentation mask (#7905) * fix: add rle compress Signed-off-by: badai-nguyen * fix: rle compress Signed-off-by: badai-nguyen * fix: move rle into utils Signed-off-by: badai-nguyen * chore: pre-commit Signed-off-by: badai-nguyen * Update perception/autoware_tensorrt_yolox/src/utils.cpp Co-authored-by: Yukihiro Saito * fix: remove unused variable Signed-off-by: badai-nguyen * Update perception/autoware_tensorrt_yolox/src/utils.cpp Co-authored-by: Manato Hirabayashi <3022416+manato@users.noreply.github.com> * style(pre-commit): autofix * feat: add unit test for utils Signed-off-by: badai-nguyen * style(pre-commit): autofix * fix: unit test Signed-off-by: badai-nguyen * chore: change to explicit index Signed-off-by: badai-nguyen * style(pre-commit): autofix * fix: cuda cmake Signed-off-by: badai-nguyen * fix: separate unit test into different PR Signed-off-by: badai-nguyen --------- Signed-off-by: badai-nguyen Co-authored-by: Yukihiro Saito Co-authored-by: Manato Hirabayashi <3022416+manato@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../autoware_tensorrt_yolox/CMakeLists.txt | 1 + .../tensorrt_yolox/tensorrt_yolox_node.hpp | 2 + .../include/autoware/tensorrt_yolox/utils.hpp | 28 ++++++++ .../src/tensorrt_yolox_node.cpp | 10 ++- .../autoware_tensorrt_yolox/src/utils.cpp | 65 +++++++++++++++++++ 5 files changed, 105 insertions(+), 1 deletion(-) create mode 100644 perception/autoware_tensorrt_yolox/include/autoware/tensorrt_yolox/utils.hpp create mode 100644 perception/autoware_tensorrt_yolox/src/utils.cpp diff --git a/perception/autoware_tensorrt_yolox/CMakeLists.txt b/perception/autoware_tensorrt_yolox/CMakeLists.txt index 1f54326f2b33a..564ac54dbc563 100644 --- 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 index 7f622b2dc867e..2e317139f27c0 100644 --- 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 #include #include +#include #include 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::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 new file mode 100644 index 0000000000000..8b86a798e798b --- /dev/null +++ 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 + +#include +#include + +namespace autoware::tensorrt_yolox +{ +std::vector> runLengthEncoder(const cv::Mat & mask); +cv::Mat runLengthDecoder(const std::vector & 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 index f3c7b2552edb6..7f398ca005d67 100644 --- 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 @@ -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> compressed_data = runLengthEncoder(mask); + int step = sizeof(uint8_t) + sizeof(int); + out_mask_msg->data.resize(static_cast(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 new file mode 100644 index 0000000000000..8d2b4a4a0527f --- /dev/null +++ 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> runLengthEncoder(const cv::Mat & image) +{ + std::vector> compressed_data; + const int rows = image.rows; + const int cols = image.cols; + compressed_data.emplace_back(image.at(0, 0), 0); + for (int i = 0; i < rows; ++i) { + for (int j = 0; j < cols; ++j) { + uint8_t current_value = image.at(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 & 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(idx / cols); + int col_idx = static_cast(idx % cols); + mask.at(row_idx, col_idx) = value; + idx++; + if (idx > rows * cols) { + break; + } + } + } + return mask; +} + +} // namespace autoware::tensorrt_yolox