Add minimal pointcloud2 support

rerun_bridge/include/rerun_bridge/rerun_ros_interface.hpp

@@ -9,6 +9,7 @@
 #include <sensor_msgs/msg/camera_info.hpp>
 #include <sensor_msgs/msg/image.hpp>
 #include <sensor_msgs/msg/imu.hpp>
+#include <sensor_msgs/msg/point_cloud2.hpp>
 #include <tf2_msgs/msg/tf_message.hpp>
 
 #include <rerun.hpp>
@@ -51,5 +52,17 @@ void log_tf_message(
 
 void log_transform(
     const rerun::RecordingStream& rec, const std::string& entity_path,
-    const geometry_msgs::msg::TransformStamped& msg
+    const geometry_msgs::msg::TransformStamped::ConstSharedPtr& msg
+);
+
+struct PointCloud2Options {
+    std::optional<std::string> color_map;
+    std::optional<std::string> color_map_field;
+    std::optional<float> color_map_min;
+    std::optional<float> color_map_max;
+};
+
+void log_point_cloud2(
+    const rerun::RecordingStream& rec, const std::string& entity_path,
+    const sensor_msgs::msg::PointCloud2::ConstSharedPtr& msg, const PointCloud2Options& options
 );

rerun_bridge/src/rerun_bridge/rerun_ros_interface.cpp

@@ -186,27 +186,71 @@ void log_camera_info(
 
 void log_transform(
     const rerun::RecordingStream& rec, const std::string& entity_path,
-    const geometry_msgs::msg::TransformStamped& msg
+    const geometry_msgs::msg::TransformStamped::ConstSharedPtr& msg
 ) {
     rec.set_time_seconds(
         "timestamp",
-        rclcpp::Time(msg.header.stamp.sec, msg.header.stamp.nanosec).seconds()
+        rclcpp::Time(msg->header.stamp.sec, msg->header.stamp.nanosec).seconds()
     );
 
     rec.log(
         entity_path,
         rerun::Transform3D(
             rerun::Vector3D(
-                msg.transform.translation.x,
-                msg.transform.translation.y,
-                msg.transform.translation.z
+                msg->transform.translation.x,
+                msg->transform.translation.y,
+                msg->transform.translation.z
             ),
             rerun::Quaternion::from_wxyz(
-                msg.transform.rotation.w,
-                msg.transform.rotation.x,
-                msg.transform.rotation.y,
-                msg.transform.rotation.z
+                msg->transform.rotation.w,
+                msg->transform.rotation.x,
+                msg->transform.rotation.y,
+                msg->transform.rotation.z
             )
         )
     );
 }
+
+void log_point_cloud2(
+    const rerun::RecordingStream& rec, const std::string& entity_path,
+    const sensor_msgs::msg::PointCloud2::ConstSharedPtr& msg, const PointCloud2Options& options
+) {
+    rec.set_time_seconds(
+        "timestamp",
+        rclcpp::Time(msg->header.stamp.sec, msg->header.stamp.nanosec).seconds()
+    );
+
+    // TODO(leo) should match the behavior described here
+    // https://wiki.ros.org/rviz/DisplayTypes/PointCloud
+    // TODO(leo) if not specified, check if 2D points or 3D points
+    // TODO(leo) allow arbitrary color mapping
+
+    size_t x_offset, y_offset, z_offset;
+    const auto bytes = msg->data;
+
+    for (const auto& field : msg->fields) {
+        if (field.name == "x") {
+            x_offset = field.offset;
+        } else if (field.name == "y") {
+            y_offset = field.offset;
+        } else if (field.name == "z") {
+            z_offset = field.offset;
+        }
+    }
+
+    std::vector<rerun::Position3D> points(msg->width * msg->height);
+
+    for (size_t i = 0; i < msg->height; ++i) {
+        for (size_t j = 0; j < msg->width; ++j) {
+            auto point_offset = i * msg->row_step + j * msg->point_step;
+            rerun::Position3D position;
+            // TODO(leo) if xyz are consecutive fields we can do this in a single memcpy
+            std::memcpy(&position.xyz.xyz[0], &bytes[point_offset + x_offset], sizeof(float));
+            std::memcpy(&position.xyz.xyz[1], &bytes[point_offset + y_offset], sizeof(float));
+            std::memcpy(&position.xyz.xyz[2], &bytes[point_offset + z_offset], sizeof(float));
+            points.emplace_back(std::move(position));
+        }
+    }
+
+    rec.log(entity_path, rerun::Points3D(points));
+}

rerun_bridge/src/rerun_bridge/visualizer_node.cpp

@@ -42,8 +42,7 @@ std::string resolve_ros_path(const std::string& path) {
 RerunLoggerNode::RerunLoggerNode() : Node("rerun_logger_node") {
     _rec.spawn().exit_on_failure();
 
-    _parallel_callback_group =
-        this->create_callback_group(rclcpp::CallbackGroupType::Reentrant);
+    _parallel_callback_group = this->create_callback_group(rclcpp::CallbackGroupType::Reentrant);
 
     _tf_buffer = std::make_unique<tf2_ros::Buffer>(this->get_clock());
     _tf_listener = std::make_unique<tf2_ros::TransformListener>(*_tf_buffer);
@@ -119,7 +118,8 @@ void RerunLoggerNode::_read_yaml_config(std::string yaml_path) {
             const std::array<float, 3> translation = {
                 extra_transform3d["transform"][3].as<float>(),
                 extra_transform3d["transform"][7].as<float>(),
-                extra_transform3d["transform"][11].as<float>()};
+                extra_transform3d["transform"][11].as<float>()
+            };
             // Rerun uses column-major order for Mat3x3
             const std::array<float, 9> mat3x3 = {
                 extra_transform3d["transform"][0].as<float>(),
@@ -130,7 +130,8 @@ void RerunLoggerNode::_read_yaml_config(std::string yaml_path) {
                 extra_transform3d["transform"][9].as<float>(),
                 extra_transform3d["transform"][2].as<float>(),
                 extra_transform3d["transform"][6].as<float>(),
-                extra_transform3d["transform"][10].as<float>()};
+                extra_transform3d["transform"][10].as<float>()
+            };
             _rec.log_timeless(
                 extra_transform3d["entity_path"].as<std::string>(),
                 rerun::Transform3D(
@@ -230,6 +231,8 @@ void RerunLoggerNode::_create_subscriptions() {
             _topic_to_subscription[topic_name] = _create_odometry_subscription(topic_name);
         } else if (topic_type == "sensor_msgs/msg/CameraInfo") {
             _topic_to_subscription[topic_name] = _create_camera_info_subscription(topic_name);
+        } else if (topic_type == "sensor_msgs/msg/PointCloud2") {
+            _topic_to_subscription[topic_name] = _create_point_cloud2_subscription(topic_name);
         }
     }
 }
@@ -257,8 +260,9 @@ void RerunLoggerNode::_update_tf() {
             continue;
         }
         try {
-            auto transform =
-                _tf_buffer->lookupTransform(parent->second, frame, now - rclcpp::Duration(1, 0));
+            auto transform = std::make_shared<geometry_msgs::msg::TransformStamped>(
+                _tf_buffer->lookupTransform(parent->second, frame, now - rclcpp::Duration(1, 0))
+            );
             log_transform(_rec, entity_path, transform);
             _last_tf_update_time = now;
         } catch (tf2::TransformException& ex) {
@@ -277,42 +281,44 @@ std::shared_ptr<rclcpp::Subscription<sensor_msgs::msg::Image>>
     RerunLoggerNode::_create_image_subscription(const std::string& topic) {
     std::string entity_path = _resolve_entity_path(topic);
     bool lookup_transform = (_topic_to_entity_path.find(topic) == _topic_to_entity_path.end());
-    ImageOptions image_options;
+    ImageOptions options;
 
     if (_topic_options.find(topic) != _topic_options.end()) {
-        image_options = _topic_options.at(topic).as<ImageOptions>();
+        options = _topic_options.at(topic).as<ImageOptions>();
     }
 
     rclcpp::SubscriptionOptions subscription_options;
     subscription_options.callback_group = _parallel_callback_group;
 
     RCLCPP_INFO(
         this->get_logger(),
-        "Subscribing to image topic %s, logging to entity path %s",
+        "Subscribing to Image topic %s, logging to entity path %s",
         topic.c_str(),
         entity_path.c_str()
     );
 
     return this->create_subscription<sensor_msgs::msg::Image>(
         topic,
         1000,
-        [&, entity_path, lookup_transform, image_options](
+        [&, entity_path, lookup_transform, options](
             const sensor_msgs::msg::Image::SharedPtr msg
         ) {
             if (!_root_frame.empty() && lookup_transform) {
                 try {
-                    auto transform = _tf_buffer->lookupTransform(
-                        _root_frame,
-                        msg->header.frame_id,
-                        msg->header.stamp,
-                        100ms
+                    auto transform = std::make_shared<geometry_msgs::msg::TransformStamped>(
+                        _tf_buffer->lookupTransform(
+                            _root_frame,
+                            msg->header.frame_id,
+                            msg->header.stamp,
+                            100ms
+                        )
                     );
                     log_transform(_rec, parent_entity_path(entity_path), transform);
                 } catch (tf2::TransformException& ex) {
                     RCLCPP_WARN(this->get_logger(), "%s", ex.what());
                 }
             }
-            log_image(_rec, entity_path, msg, image_options);
+            log_image(_rec, entity_path, msg, options);
         },
         subscription_options
     );
@@ -405,6 +411,60 @@ std::shared_ptr<rclcpp::Subscription<sensor_msgs::msg::CameraInfo>>
     );
 }
 
+std::shared_ptr<rclcpp::Subscription<sensor_msgs::msg::PointCloud2>>
+    RerunLoggerNode::_create_point_cloud2_subscription(const std::string& topic) {
+    std::string entity_path = _resolve_entity_path(topic);
+    bool lookup_transform = (_topic_to_entity_path.find(topic) == _topic_to_entity_path.end());
+    bool restamp_msgs = true;  // TODO make this configurable and applicable to all types
+    PointCloud2Options options;
+
+    if (_topic_options.find(topic) != _topic_options.end()) {
+        options = _topic_options.at(topic).as<PointCloud2Options>();
+    }
+
+    rclcpp::SubscriptionOptions subscription_options;
+    subscription_options.callback_group = _parallel_callback_group;
+
+    RCLCPP_INFO(
+        this->get_logger(),
+        "Subscribing to PointCloud2 topic %s, logging to entity path %s",
+        topic.c_str(),
+        entity_path.c_str()
+    );
+
+    return this->create_subscription<sensor_msgs::msg::PointCloud2>(
+        topic,
+        1000,
+        [&, entity_path, lookup_transform, options, restamp_msgs](
+            const sensor_msgs::msg::PointCloud2::SharedPtr msg
+        ) {
+            if(restamp_msgs) {
+                auto now = this->get_clock()->now();
+                msg->header.stamp = now;
+            }
+
+            if (!_root_frame.empty() && lookup_transform) {
+                try {
+                    auto transform = std::make_shared<geometry_msgs::msg::TransformStamped>(
+                        _tf_buffer->lookupTransform(
+                            _root_frame,
+                            msg->header.frame_id,
+                            msg->header.stamp,
+                            100ms
+                        )
+                    );
+                    log_transform(_rec, parent_entity_path(entity_path), transform);
+                } catch (tf2::TransformException& ex) {
+                    RCLCPP_WARN(this->get_logger(), "%s", ex.what());
+                }
+            }
+
+            log_point_cloud2(_rec, entity_path, msg, options);
+        },
+        subscription_options
+    );
+}
+
 namespace YAML {
     template <>
     struct convert<ImageOptions> {
@@ -431,6 +491,40 @@ namespace YAML {
             return true;
         }
     };
+
+    template <>
+    struct convert<PointCloud2Options> {
+        static bool decode(const Node& node, PointCloud2Options& rhs) {
+            int total = 0;
+
+            if (!node.IsMap()) {
+                return false;
+            }
+
+            if (node["color_map"]) {
+                rhs.color_map = node["color_map"].as<std::string>();
+                ++total;
+            }
+            if (node["color_map_field"]) {
+                rhs.color_map_field = node["color_map_field"].as<std::string>();
+                ++total;
+            }
+            if (node["color_map_min"]) {
+                rhs.color_map_min = node["color_map_min"].as<float>();
+                ++total;
+            }
+            if (node["color_map_max"]) {
+                rhs.color_map_max = node["color_map_max"].as<float>();
+                ++total;
+            }
+
+            if (total != node.size()) {
+                return false;
+            }
+
+            return true;
+        }
+    };
 } // namespace YAML
 
 int main(int argc, char** argv) {

rerun_bridge/src/rerun_bridge/visualizer_node.hpp

@@ -13,6 +13,7 @@
 #include <sensor_msgs/msg/camera_info.hpp>
 #include <sensor_msgs/msg/image.hpp>
 #include <sensor_msgs/msg/imu.hpp>
+#include <sensor_msgs/msg/point_cloud2.hpp>
 #include <tf2_msgs/msg/tf_message.hpp>
 
 #include <rerun.hpp>
@@ -69,4 +70,6 @@ class RerunLoggerNode : public rclcpp::Node {
     );
     std::shared_ptr<rclcpp::Subscription<sensor_msgs::msg::CameraInfo>>
         _create_camera_info_subscription(const std::string& topic);
+    std::shared_ptr<rclcpp::Subscription<sensor_msgs::msg::PointCloud2>>
+        _create_point_cloud2_subscription(const std::string& topic);
 };