Extend and refactor log_point_cloud2

* log_point_cloud2 refactored to avoid redundant allocations * PointCloudProcessor and ColorMapper added * reversed colormap added * ColormapsLUT::supportedColormaps added * additional compilation flags added * rerun-sdk verison updated
rerun-io · Aug 7, 2024 · 4fae4d5 · 4fae4d5
1 parent bc8cb32
commit 4fae4d5
Show file tree

Hide file tree

Showing 11 changed files with 448 additions and 212 deletions.
diff --git a/rerun_bridge/CMakeLists.txt b/rerun_bridge/CMakeLists.txt
@@ -4,6 +4,7 @@ project(rerun_bridge)
 if(NOT DEFINED CMAKE_CXX_STANDARD)
   set(CMAKE_CXX_STANDARD 17)
 endif()
+add_compile_options(-Wall -Wextra -Wpedantic)
 
 # Avoid warning about CMP0135
 if(CMAKE_VERSION VERSION_GREATER_EQUAL "3.24.0")
@@ -22,11 +23,16 @@ find_package(OpenCV REQUIRED)
 find_package(yaml-cpp REQUIRED)
 
 include(FetchContent)
-FetchContent_Declare(rerun_sdk URL https://github.com/rerun-io/rerun/releases/download/0.16.0/rerun_cpp_sdk.zip)
+FetchContent_Declare(rerun_sdk URL https://github.com/rerun-io/rerun/releases/download/0.17.0/rerun_cpp_sdk.zip)
 FetchContent_MakeAvailable(rerun_sdk)
 
 # setup targets (has to be done before ament_package call)
-add_library(${PROJECT_NAME} src/rerun_bridge/rerun_ros_interface.cpp)
+add_library(${PROJECT_NAME} 
+  src/rerun_bridge/rerun_ros_interface.cpp
+  src/rerun_bridge/color_maps.cpp
+  src/rerun_bridge/point_cloud_processor.cpp
+  src/rerun_bridge/color_mapper.cpp
+)
 add_executable(visualizer src/rerun_bridge/visualizer_node.cpp)
 
 # add system dependencies

diff --git a/rerun_bridge/include/rerun_bridge/color_mapper.hpp b/rerun_bridge/include/rerun_bridge/color_mapper.hpp
@@ -0,0 +1,44 @@
+#pragma once
+
+#include <optional>
+#include <string>
+#include <vector>
+
+#include "sensor_msgs/point_cloud2_iterator.hpp"
+
+#include "rerun_bridge/color_maps.hpp"
+#include "rerun_bridge/rerun_ros_interface.hpp"
+
+using pclConstIter = sensor_msgs::PointCloud2ConstIterator<float>;
+
+struct PointCloud2Options {
+    std::optional<std::string> colormapName;
+    std::optional<std::string> colormapField;
+    std::optional<float> colormapMin;
+    std::optional<float> colormapMax;
+};
+
+class ColorMapper {
+    std::optional<colormapLUT> colormap_{};
+    std::vector<float> fieldValues_{};
+    std::vector<rerun::Color> colors_{};
+
+    std::pair<float, float> getMinMax(const std::vector<float>& values) const;
+    std::size_t calculateIdx(const float& value, const float& minValue, const float& maxValue)
+        const;
+    void setColormap();
+    std::string toUppercase(std::string str) const;
+    void convertPclMsgToColorapFieldVec(const sensor_msgs::msg::PointCloud2::ConstSharedPtr& msg);
+    void remapValuesToColors();
+
+  public:
+    // data from yaml file
+    const PointCloud2Options options_{};
+
+    ColorMapper(const PointCloud2Options& options);
+    void calculateRerunColors(const sensor_msgs::msg::PointCloud2::ConstSharedPtr& msg);
+    bool isValid() const;
+    void reserve(const size_t& size);
+    void clear();
+    const std::vector<rerun::components::Color>& getColors() const;
+};
diff --git a/rerun_bridge/include/rerun_bridge/color_maps.hpp b/rerun_bridge/include/rerun_bridge/color_maps.hpp
@@ -0,0 +1,13 @@
+#pragma once
+
+#include <array>
+#include <unordered_map>
+
+using colormapLUT = std::array<std::array<float, 3>, 256>;
+
+namespace ColormapsLUT {
+    extern const std::unordered_map<std::string, colormapLUT> supportedColormaps;
+
+    extern const colormapLUT Turbo;
+    extern const colormapLUT Rainbow;
+} // namespace ColormapsLUT
diff --git a/rerun_bridge/include/rerun_bridge/point_cloud_processor.hpp b/rerun_bridge/include/rerun_bridge/point_cloud_processor.hpp
@@ -0,0 +1,46 @@
+#pragma once
+
+#include <algorithm>
+#include <optional>
+#include <vector>
+
+#include "rerun.hpp"
+#include "rerun_bridge/color_mapper.hpp" // PointCloud2Options
+#include "rerun_bridge/rerun_ros_interface.hpp"
+
+class PointCloudProcessor {
+    ColorMapper colorMapper_;
+    std::size_t maxNumPerMsg_{};
+    std::vector<rerun::Position3D> points_{};
+
+    std::pair<std::optional<float>, std::optional<float>> getMinMax(const std::vector<float>& values
+    ) const;
+
+    void reserve(const sensor_msgs::msg::PointCloud2::ConstSharedPtr& msg);
+    bool isMsgLonger(const sensor_msgs::msg::PointCloud2::ConstSharedPtr& msg) const;
+    const std::vector<rerun::components::Position3D>& convertPclMsgToPosition3DVec(
+        const sensor_msgs::msg::PointCloud2::ConstSharedPtr& msg
+    );
+    void convertPclMsgToColorMapFieldVec(
+        const sensor_msgs::msg::PointCloud2::ConstSharedPtr& msg, std::vector<float>& values
+    ) const;
+    void convertPclMsgToPosition3DVecIter(
+        const sensor_msgs::msg::PointCloud2::ConstSharedPtr& msg,
+        std::vector<rerun::components::Position3D>& points
+    ) const;
+    bool areFieldNamesSupported(
+        const rerun::RecordingStream& rec, const std::string& entity_path,
+        const sensor_msgs::msg::PointCloud2::ConstSharedPtr& msg
+    ) const;
+    bool isFieldTypeFloat(
+        const rerun::RecordingStream& rec, const std::string& entity_path,
+        const sensor_msgs::msg::PointField& field
+    ) const;
+
+  public:
+    PointCloudProcessor(const PointCloud2Options& options);
+    void logPointCloud2(
+        const rerun::RecordingStream& rec, const std::string& entity_path,
+        const sensor_msgs::msg::PointCloud2::ConstSharedPtr& msg
+    );
+};
diff --git a/rerun_bridge/include/rerun_bridge/rerun_ros_interface.hpp b/rerun_bridge/include/rerun_bridge/rerun_ros_interface.hpp
@@ -1,6 +1,7 @@
 #pragma once
 
 #include <map>
+#include <optional>
 #include <string>
 
 #include <geometry_msgs/msg/pose_stamped.hpp>
@@ -14,16 +15,16 @@
 
 #include <rerun.hpp>
 
+struct ImageOptions {
+    std::optional<float> min_depth{};
+    std::optional<float> max_depth{};
+};
+
 void log_imu(
     const rerun::RecordingStream& rec, const std::string& entity_path,
     const sensor_msgs::msg::Imu::ConstSharedPtr& msg
 );
 
-struct ImageOptions {
-    std::optional<float> min_depth;
-    std::optional<float> max_depth;
-};
-
 void log_image(
     const rerun::RecordingStream& rec, const std::string& entity_path,
     const sensor_msgs::msg::Image::ConstSharedPtr& msg, const ImageOptions& options = ImageOptions{}
@@ -54,15 +55,3 @@ void log_transform(
     const rerun::RecordingStream& rec, const std::string& entity_path,
     const geometry_msgs::msg::TransformStamped::ConstSharedPtr& msg
 );
-
-struct PointCloud2Options {
-    std::optional<std::string> colormap;
-    std::optional<std::string> colormap_field;
-    std::optional<float> colormap_min;
-    std::optional<float> colormap_max;
-};
-
-void log_point_cloud2(
-    const rerun::RecordingStream& rec, const std::string& entity_path,
-    const sensor_msgs::msg::PointCloud2::ConstSharedPtr& msg, const PointCloud2Options& options
-);
diff --git a/rerun_bridge/src/rerun_bridge/color_mapper.cpp b/rerun_bridge/src/rerun_bridge/color_mapper.cpp
@@ -0,0 +1,100 @@
+#include <algorithm>
+#include <cctype> // std::toupper
+#include <iostream>
+#include <optional>
+#include <string>
+
+#include "rerun_bridge/color_mapper.hpp"
+#include "rerun_bridge/color_maps.hpp"
+
+ColorMapper::ColorMapper(const PointCloud2Options& options) : options_{options} {
+    setColormap();
+}
+
+bool ColorMapper::isValid() const {
+    return (colormap_.has_value() && options_.colormapField.has_value());
+}
+
+void ColorMapper::calculateRerunColors(const sensor_msgs::msg::PointCloud2::ConstSharedPtr& msg) {
+    convertPclMsgToColorapFieldVec(msg);
+    remapValuesToColors();
+}
+
+std::size_t ColorMapper::calculateIdx(
+    const float& value, const float& minValue, const float& maxValue
+) const {
+    return static_cast<size_t>(255 * (value - minValue) / (maxValue - minValue));
+}
+
+std::pair<float, float> ColorMapper::getMinMax(const std::vector<float>& values) const {
+    float min_value{0}, max_value{0};
+    if (!options_.colormapMin) {
+        min_value = *std::min_element(values.begin(), values.end());
+    } else {
+        min_value = *options_.colormapMin;
+    }
+
+    if (!options_.colormapMax) {
+        max_value = *std::max_element(values.begin(), values.end());
+    } else {
+        max_value = *options_.colormapMax;
+    }
+    return {min_value, max_value};
+}
+
+void ColorMapper::reserve(const size_t& size) {
+    colors_.reserve(size);
+    fieldValues_.reserve(size);
+}
+
+void ColorMapper::clear() {
+    colors_.clear();
+    fieldValues_.clear();
+}
+
+std::string ColorMapper::toUppercase(std::string str) const {
+    std::transform(str.begin(), str.end(), str.begin(), [](unsigned char c) {
+        return std::toupper(c);
+    });
+    return str;
+}
+
+void ColorMapper::setColormap() {
+    if (!options_.colormapName.has_value()) {
+        return;
+    }
+
+    // name from YAML file
+    auto colormapName = toUppercase(*options_.colormapName);
+    auto it = ColormapsLUT::supportedColormaps.find(colormapName);
+    if (it != ColormapsLUT::supportedColormaps.end()) {
+        colormap_ = it->second;
+    } else {
+        std::cout << "Colormap: " << colormapName << " is not supported" << std::endl;
+    }
+}
+
+void ColorMapper::convertPclMsgToColorapFieldVec(
+    const sensor_msgs::msg::PointCloud2::ConstSharedPtr& msg
+) {
+    auto iterColorField = pclConstIter(*msg, options_.colormapField.value());
+    for (; iterColorField != iterColorField.end(); ++iterColorField) {
+        fieldValues_.emplace_back(*iterColorField);
+    }
+}
+
+void ColorMapper::remapValuesToColors() {
+    if (!options_.colormapName || !colormap_) {
+        return;
+    }
+
+    const auto [min, max] = getMinMax(fieldValues_);
+    for (const auto& value : fieldValues_) {
+        const auto idx = calculateIdx(value, min, max);
+        colors_.emplace_back((*colormap_)[idx][0], (*colormap_)[idx][1], (*colormap_)[idx][2]);
+    }
+}
+
+const std::vector<rerun::components::Color>& ColorMapper::getColors() const {
+    return colors_;
+}
diff --git a/rerun_bridge/src/rerun_bridge/color_maps.cpp b/rerun_bridge/src/rerun_bridge/color_maps.cpp
@@ -0,0 +1,84 @@
+#include <array>
+#include <string>
+#include <tuple>
+#include <unordered_map>
+
+#include "rerun_bridge/color_maps.hpp"
+
+constexpr colormapLUT reverse(const colormapLUT& input) {
+    colormapLUT reversed{};
+    constexpr std::size_t N = std::tuple_size<colormapLUT>::value;
+    for (std::size_t i = 0; i < N; ++i) {
+        reversed[i] = input[N - 1 - i];
+    }
+    return reversed;
+}
+
+namespace ColormapsLUT {
+    const std::unordered_map<std::string, colormapLUT> supportedColormaps = {
+        {"TURBO", Turbo}, {"RAINBOW", Rainbow}};
+
+    // Turbo colormap lookup table
+    // from: https://gist.github.com/mikhailov-work/6a308c20e494d9e0ccc29036b28faa7a
+    // Copyright 2019 Google LLC.
+    // SPDX-License-Identifier: Apache-2.0
+    //
+    // Author: Anton Mikhailov
+
+    constexpr colormapLUT Turbo = {
+        {{48, 18, 59},    {50, 21, 67},    {51, 24, 74},   {52, 27, 81},   {53, 30, 88},
+         {54, 33, 95},    {55, 36, 102},   {56, 39, 109},  {57, 42, 115},  {58, 45, 121},
+         {59, 47, 128},   {60, 50, 134},   {61, 53, 139},  {62, 56, 145},  {63, 59, 151},
+         {63, 62, 156},   {64, 64, 162},   {65, 67, 167},  {65, 70, 172},  {66, 73, 177},
+         {66, 75, 181},   {67, 78, 186},   {68, 81, 191},  {68, 84, 195},  {68, 86, 199},
+         {69, 89, 203},   {69, 92, 207},   {69, 94, 211},  {70, 97, 214},  {70, 100, 218},
+         {70, 102, 221},  {70, 105, 224},  {70, 107, 227}, {71, 110, 230}, {71, 113, 233},
+         {71, 115, 235},  {71, 118, 238},  {71, 120, 240}, {71, 123, 242}, {70, 125, 244},
+         {70, 128, 246},  {70, 130, 248},  {70, 133, 250}, {70, 135, 251}, {69, 138, 252},
+         {69, 140, 253},  {68, 143, 254},  {67, 145, 254}, {66, 148, 255}, {65, 150, 255},
+         {64, 153, 255},  {62, 155, 254},  {61, 158, 254}, {59, 160, 253}, {58, 163, 252},
+         {56, 165, 251},  {55, 168, 250},  {53, 171, 248}, {51, 173, 247}, {49, 175, 245},
+         {47, 178, 244},  {46, 180, 242},  {44, 183, 240}, {42, 185, 238}, {40, 188, 235},
+         {39, 190, 233},  {37, 192, 231},  {35, 195, 228}, {34, 197, 226}, {32, 199, 223},
+         {31, 201, 221},  {30, 203, 218},  {28, 205, 216}, {27, 208, 213}, {26, 210, 210},
+         {26, 212, 208},  {25, 213, 205},  {24, 215, 202}, {24, 217, 200}, {24, 219, 197},
+         {24, 221, 194},  {24, 222, 192},  {24, 224, 189}, {25, 226, 187}, {25, 227, 185},
+         {26, 228, 182},  {28, 230, 180},  {29, 231, 178}, {31, 233, 175}, {32, 234, 172},
+         {34, 235, 170},  {37, 236, 167},  {39, 238, 164}, {42, 239, 161}, {44, 240, 158},
+         {47, 241, 155},  {50, 242, 152},  {53, 243, 148}, {56, 244, 145}, {60, 245, 142},
+         {63, 246, 138},  {67, 247, 135},  {70, 248, 132}, {74, 248, 128}, {78, 249, 125},
+         {82, 250, 122},  {85, 250, 118},  {89, 251, 115}, {93, 252, 111}, {97, 252, 108},
+         {101, 253, 105}, {105, 253, 102}, {109, 254, 98}, {113, 254, 95}, {117, 254, 92},
+         {121, 254, 89},  {125, 255, 86},  {128, 255, 83}, {132, 255, 81}, {136, 255, 78},
+         {139, 255, 75},  {143, 255, 73},  {146, 255, 71}, {150, 254, 68}, {153, 254, 66},
+         {156, 254, 64},  {159, 253, 63},  {161, 253, 61}, {164, 252, 60}, {167, 252, 58},
+         {169, 251, 57},  {172, 251, 56},  {175, 250, 55}, {177, 249, 54}, {180, 248, 54},
+         {183, 247, 53},  {185, 246, 53},  {188, 245, 52}, {190, 244, 52}, {193, 243, 52},
+         {195, 241, 52},  {198, 240, 52},  {200, 239, 52}, {203, 237, 52}, {205, 236, 52},
+         {208, 234, 52},  {210, 233, 53},  {212, 231, 53}, {215, 229, 53}, {217, 228, 54},
+         {219, 226, 54},  {221, 224, 55},  {223, 223, 55}, {225, 221, 55}, {227, 219, 56},
+         {229, 217, 56},  {231, 215, 57},  {233, 213, 57}, {235, 211, 57}, {236, 209, 58},
+         {238, 207, 58},  {239, 205, 58},  {241, 203, 58}, {242, 201, 58}, {244, 199, 58},
+         {245, 197, 58},  {246, 195, 58},  {247, 193, 58}, {248, 190, 57}, {249, 188, 57},
+         {250, 186, 57},  {251, 184, 56},  {251, 182, 55}, {252, 179, 54}, {252, 177, 54},
+         {253, 174, 53},  {253, 172, 52},  {254, 169, 51}, {254, 167, 50}, {254, 164, 49},
+         {254, 161, 48},  {254, 158, 47},  {254, 155, 45}, {254, 153, 44}, {254, 150, 43},
+         {254, 147, 42},  {254, 144, 41},  {253, 141, 39}, {253, 138, 38}, {252, 135, 37},
+         {252, 132, 35},  {251, 129, 34},  {251, 126, 33}, {250, 123, 31}, {249, 120, 30},
+         {249, 117, 29},  {248, 114, 28},  {247, 111, 26}, {246, 108, 25}, {245, 105, 24},
+         {244, 102, 23},  {243, 99, 21},   {242, 96, 20},  {241, 93, 19},  {240, 91, 18},
+         {239, 88, 17},   {237, 85, 16},   {236, 83, 15},  {235, 80, 14},  {234, 78, 13},
+         {232, 75, 12},   {231, 73, 12},   {229, 71, 11},  {228, 69, 10},  {226, 67, 10},
+         {225, 65, 9},    {223, 63, 8},    {221, 61, 8},   {220, 59, 7},   {218, 57, 7},
+         {216, 55, 6},    {214, 53, 6},    {212, 51, 5},   {210, 49, 5},   {208, 47, 5},
+         {206, 45, 4},    {204, 43, 4},    {202, 42, 4},   {200, 40, 3},   {197, 38, 3},
+         {195, 37, 3},    {193, 35, 2},    {190, 33, 2},   {188, 32, 2},   {185, 30, 2},
+         {183, 29, 2},    {180, 27, 1},    {178, 26, 1},   {175, 24, 1},   {172, 23, 1},
+         {169, 22, 1},    {167, 20, 1},    {164, 19, 1},   {161, 18, 1},   {158, 16, 1},
+         {155, 15, 1},    {152, 14, 1},    {149, 13, 1},   {146, 11, 1},   {142, 10, 1},
+         {139, 9, 2},     {136, 8, 2},     {133, 7, 2},    {129, 6, 2},    {126, 5, 2},
+         {122, 4, 3}}};
+
+    constexpr colormapLUT Rainbow = reverse(Turbo);
+
+} // namespace ColormapsLUT