Feature/better object viz (#38)

nathanhhughes · nathanhhughes · commit 05a6b00949e4 · 2025-09-30T13:07:36.000-04:00
* publish all labels in same array

* clean up handling of visualization images w.r.t. sensors

* use display config for tracking image also
diff --git a/hydra_ros/include/hydra_ros/active_window/reconstruction_visualizer.h b/hydra_ros/include/hydra_ros/active_window/reconstruction_visualizer.h
@@ -34,6 +34,7 @@
  * -------------------------------------------------------------------------- */
 #pragma once
 #include <hydra/active_window/active_window_module.h>
+#include <hydra/input/sensor_map.h>
 #include <hydra_visualizer/adapters/mesh_color.h>
 #include <hydra_visualizer/color/colormap_utilities.h>
 #include <hydra_visualizer/utils/marker_group_pub.h>
@@ -51,6 +52,12 @@ namespace hydra {
 
 class ReconstructionVisualizer : public ActiveWindowModule::Sink {
  public:
+  struct SensorDisplay {
+    using Config = DisplayConfig;
+    const Config config;
+    explicit SensorDisplay(const Config& config) : config(config) {}
+  };
+
   struct Config {
     std::string ns = "~/reconstruction";
     double min_weight = 0.0;
@@ -70,7 +77,7 @@ class ReconstructionVisualizer : public ActiveWindowModule::Sink {
     visualizer::RangeColormap::Config colormap;
     visualizer::CategoricalColormap::Config label_colormap;
     config::VirtualConfig<MeshColoring> mesh_coloring;
-    DisplayConfig image_display;
+    SensorMap<SensorDisplay>::Config sensor_displays;
   } const config;
 
   explicit ReconstructionVisualizer(const Config& config);
@@ -90,6 +97,7 @@ class ReconstructionVisualizer : public ActiveWindowModule::Sink {
   MarkerGroupPub pubs_;
   rclcpp::Publisher<kimera_pgmo_msgs::msg::Mesh>::SharedPtr active_mesh_pub_;
   rclcpp::Publisher<geometry_msgs::msg::PoseStamped>::SharedPtr pose_pub_;
+  SensorMap<SensorDisplay> sensor_displays_;
   ianvs::RosPublisherGroup<sensor_msgs::msg::Image> image_pubs_;
   ianvs::RosPublisherGroup<sensor_msgs::msg::PointCloud2> cloud_pubs_;
   const visualizer::RangeColormap colormap_;
diff --git a/hydra_ros/include/hydra_ros/frontend/object_visualizer.h b/hydra_ros/include/hydra_ros/frontend/object_visualizer.h
@@ -37,7 +37,7 @@
 #include <ianvs/lazy_publisher_group.h>
 #include <kimera_pgmo/mesh_delta.h>
 
-#include <visualization_msgs/msg/marker.hpp>
+#include <visualization_msgs/msg/marker_array.hpp>
 
 namespace hydra {
 
@@ -68,7 +68,7 @@ class ObjectVisualizer : public MeshSegmenter::Sink {
 
  protected:
   ianvs::NodeHandle nh_;
-  ianvs::RosPublisherGroup<visualization_msgs::msg::Marker> pubs_;
+  ianvs::RosPublisherGroup<visualization_msgs::msg::MarkerArray> pubs_;
 
  private:
   inline static const auto registration_ =
diff --git a/hydra_ros/include/hydra_ros/utils/input_data_to_messages.h b/hydra_ros/include/hydra_ros/utils/input_data_to_messages.h
@@ -33,46 +33,59 @@
  * purposes notwithstanding any copyright notation herein.
  * -------------------------------------------------------------------------- */
 #pragma once
+#include <hydra_visualizer/color/colormap_utilities.h>
 #include <spark_dsg/color.h>
 
 #include <functional>
+#include <limits>
 
+#include <opencv2/core/mat.hpp>
 #include <sensor_msgs/msg/image.hpp>
 #include <sensor_msgs/msg/point_cloud2.hpp>
 
 namespace hydra {
 
 struct InputData;
-using CmapFunc = std::function<spark_dsg::Color(uint32_t)>;
+using CmapFunc = std::function<spark_dsg::Color(cv::Mat, int, int)>;
 
 struct DisplayConfig {
   float width_scale = 1.0f;
   float height_scale = 1.0f;
+  float overlay_alpha = 0.5f;
+  float min_distance = -1.0f;
+  float max_distance = -1.0f;
+  visualizer::RangeColormap::Config distance_colormap;
 };
 
 void declare_config(DisplayConfig& config);
 
+/**
+ * @brief Turn an image into a message, optionally applying display config
+ */
+sensor_msgs::msg::Image::SharedPtr convertImage(const std_msgs::msg::Header& header,
+                                                const cv::Mat& img,
+                                                const DisplayConfig& config);
+
 /**
  * @brief Make a colored image for the current labels in the input data
  */
 sensor_msgs::msg::Image::SharedPtr makeImage(const std_msgs::msg::Header& header,
-                                             const InputData& sensor_data,
+                                             const cv::Mat& img_in,
                                              const CmapFunc& colormap,
                                              const DisplayConfig& config = {});
 
 /**
- * @brief Copy the current depth image to ros
+ * @brief Make a colored image for the current labels in the input data
  */
-sensor_msgs::msg::Image::SharedPtr makeDepthImage(const std_msgs::msg::Header& header,
-                                                  const InputData& sensor_data,
-                                                  const DisplayConfig& config = {});
+sensor_msgs::msg::Image::SharedPtr makeOverlayImage(const std_msgs::msg::Header& header,
+                                                    const cv::Mat& img_in,
+                                                    const cv::Mat& color_in,
+                                                    const CmapFunc& colormap,
+                                                    const DisplayConfig& config = {});
 
-/**
- * @brief Copy the current range image to ros
- */
-sensor_msgs::msg::Image::SharedPtr makeRangeImage(const std_msgs::msg::Header& header,
-                                                  const InputData& sensor_data,
-                                                  const DisplayConfig& config = {});
+sensor_msgs::msg::Image::SharedPtr makeDistImage(const std_msgs::msg::Header& header,
+                                                 const cv::Mat& img_in,
+                                                 const DisplayConfig& config = {});
 
 /**
  * @brief Convert the input pointcloud to a ROS type
diff --git a/hydra_ros/src/active_window/reconstruction_visualizer.cpp b/hydra_ros/src/active_window/reconstruction_visualizer.cpp
@@ -99,7 +99,7 @@ void declare_config(ReconstructionVisualizer::Config& config) {
   field(config.filter_points_by_range, "filter_points_by_range");
   field(config.colormap, "colormap");
   field(config.label_colormap, "label_colormap");
-  field(config.image_display, "image_display");
+  field(config.sensor_displays, "sensor_displays");
   config.mesh_coloring.setOptional();
   field(config.mesh_coloring, "mesh_coloring");
 }
@@ -110,6 +110,7 @@ ReconstructionVisualizer::ReconstructionVisualizer(const Config& config)
       pubs_(nh_),
       active_mesh_pub_(nh_.create_publisher<kimera_pgmo_msgs::msg::Mesh>("mesh", 1)),
       pose_pub_(nh_.create_publisher<geometry_msgs::msg::PoseStamped>("pose", 10)),
+      sensor_displays_(config.sensor_displays),
       image_pubs_(nh_),
       cloud_pubs_(nh_),
       colormap_(config.colormap),
@@ -182,24 +183,36 @@ void ReconstructionVisualizer::call(uint64_t timestamp_ns,
     std_msgs::msg::Header header;
     header.stamp = rclcpp::Time(output.timestamp_ns);
     header.frame_id = GlobalInfo::instance().getFrames().map;
-
     const auto sensor_name = output.sensor_data->getSensor().name;
+
+    DisplayConfig display_config;
+    const auto display = sensor_displays_.get(sensor_name);
+    if (display) {
+      display_config = display->config;
+    }
+
     image_pubs_.publish(sensor_name + "/labels", [&]() {
-      return makeImage(
+      return makeOverlayImage(
           header,
-          *output.sensor_data,
-          [this](uint32_t label) { return label_colormap_(label); },
-          config.image_display);
-    });
-    image_pubs_.publish(sensor_name + "/depth", [&]() {
-      return makeDepthImage(header, *output.sensor_data, config.image_display);
+          output.sensor_data->label_image,
+          output.sensor_data->color_image,
+          [this](const cv::Mat& img, int r, int c) {
+            return label_colormap_(img.at<int32_t>(r, c));
+          },
+          display_config);
     });
     image_pubs_.publish(sensor_name + "/range", [&]() {
-      return makeRangeImage(header, *output.sensor_data, config.image_display);
+      return makeDistImage(header, output.sensor_data->range_image, display_config);
     });
     cloud_pubs_.publish(sensor_name + "/pointcloud", [&]() {
       return makeCloud(header, *output.sensor_data, config.filter_points_by_range);
     });
+
+    if (!output.sensor_data->depth_image.empty()) {
+      image_pubs_.publish(sensor_name + "/depth", [&]() {
+        return makeDistImage(header, output.sensor_data->depth_image, display_config);
+      });
+    }
   }
 }  // namespace hydra
 
diff --git a/hydra_ros/src/frontend/object_visualizer.cpp b/hydra_ros/src/frontend/object_visualizer.cpp
@@ -42,6 +42,7 @@
 namespace hydra {
 
 using visualization_msgs::msg::Marker;
+using visualization_msgs::msg::MarkerArray;
 
 void declare_config(ObjectVisualizer::Config& config) {
   using namespace config;
@@ -61,30 +62,21 @@ std::string ObjectVisualizer::printInfo() const { return config::toString(config
 
 void ObjectVisualizer::call(uint64_t timestamp_ns,
                             const kimera_pgmo::MeshDelta& delta,
-                            const std::vector<size_t>& active,
+                            const std::vector<size_t>& /* active */,
                             const LabelIndices& label_indices) const {
-  pubs_.publish("active_vertices", [&]() {
-    auto msg = std::make_unique<Marker>();
-    msg->header.stamp = rclcpp::Time(timestamp_ns);
-    msg->header.frame_id = GlobalInfo::instance().getFrames().odom;
-    msg->ns = "active_vertices";
-    msg->id = 0;
-    fillMarkerFromCloud(delta, active, *msg);
-    return msg;
-  });
+  pubs_.publish("object_vertices", [&]() {
+    auto msg_arr = std::make_unique<MarkerArray>();
+    for (const auto& [label, indices] : label_indices) {
+      auto& msg = msg_arr->markers.emplace_back();
+      msg.header.stamp = rclcpp::Time(timestamp_ns);
+      msg.header.frame_id = GlobalInfo::instance().getFrames().odom;
+      msg.ns = "label_vertices_" + std::to_string(label);
+      msg.id = 0;
+      fillMarkerFromCloud(delta, indices, msg);
+    }
 
-  for (const auto& id_label_pair : label_indices) {
-    pubs_.publish("object_vertices/label" + std::to_string(id_label_pair.first), [&]() {
-      const auto& [label, indices] = id_label_pair;
-      auto msg = std::make_unique<Marker>();
-      msg->header.stamp = rclcpp::Time(timestamp_ns);
-      msg->header.frame_id = GlobalInfo::instance().getFrames().odom;
-      msg->ns = "label_vertices_" + std::to_string(label);
-      msg->id = 0;
-      fillMarkerFromCloud(delta, indices, *msg);
-      return msg;
-    });
-  }
+    return msg_arr;
+  });
 }
 
 void ObjectVisualizer::fillMarkerFromCloud(const kimera_pgmo::MeshDelta& delta,
diff --git a/hydra_ros/src/utils/input_data_to_messages.cpp b/hydra_ros/src/utils/input_data_to_messages.cpp
@@ -60,7 +60,7 @@ PointField makeField(const std::string& name,
   return field;
 }
 
-cv::Mat showImage(const cv::Mat& input, const DisplayConfig& config) {
+cv::Mat resizeImage(const cv::Mat& input, const DisplayConfig& config) {
   if (config.width_scale == 1.0f && config.height_scale == 1.0f) {
     return input;
   }
@@ -82,49 +82,85 @@ void declare_config(DisplayConfig& config) {
   name("DisplayConfig");
   field(config.width_scale, "width_scale");
   field(config.height_scale, "height_scale");
+  field(config.overlay_alpha, "overlay_alpha");
+  field(config.min_distance, "min_distance", "m");
+  field(config.max_distance, "max_distance", "m");
+  field(config.distance_colormap, "distance_colormap", "m");
+  // overlay == 0.0 doesn't make sense, but overlay == 1.0 does
+  checkInRange(config.overlay_alpha, 0.0f, 1.0f, "overlay_alpha", false, true);
 }
 
-Image::SharedPtr makeImage(const std_msgs::msg::Header& header,
-                           const InputData& sensor_data,
-                           const CmapFunc& colormap,
-                           const DisplayConfig& config) {
-  const auto& labels = sensor_data.label_image;
+Image::SharedPtr convertImage(const std_msgs::msg::Header& header,
+                              const cv::Mat& img,
+                              const DisplayConfig& config) {
   cv_bridge::CvImagePtr msg(new cv_bridge::CvImage());
   msg->header = header;
   msg->encoding = "rgb8";
-  cv::Mat img(labels.rows, labels.cols, CV_8UC3);
-  for (int r = 0; r < labels.rows; ++r) {
-    for (int c = 0; c < labels.cols; ++c) {
+  msg->image = resizeImage(img, config);
+  return msg->toImageMsg();
+}
+
+Image::SharedPtr makeImage(const std_msgs::msg::Header& header,
+                           const cv::Mat& img_in,
+                           const CmapFunc& colormap,
+                           const DisplayConfig& config) {
+  cv::Mat img(img_in.rows, img_in.cols, CV_8UC3);
+  for (int r = 0; r < img_in.rows; ++r) {
+    for (int c = 0; c < img_in.cols; ++c) {
       auto pixel = img.ptr<uint8_t>(r, c);
-      const auto color = colormap(labels.at<int>(r, c));
+      const auto color = colormap(img_in, r, c);
       *pixel = color.r;
       *(pixel + 1) = color.g;
       *(pixel + 2) = color.b;
     }
   }
 
-  msg->image = showImage(img, config);
-  return msg->toImageMsg();
+  return convertImage(header, img, config);
 }
 
-Image::SharedPtr makeDepthImage(const std_msgs::msg::Header& header,
-                                const InputData& sensor_data,
-                                const DisplayConfig& config) {
-  cv_bridge::CvImagePtr msg(new cv_bridge::CvImage());
-  msg->header = header;
-  msg->encoding = "32FC1";
-  msg->image = showImage(sensor_data.depth_image, config);
-  return msg->toImageMsg();
+Image::SharedPtr makeOverlayImage(const std_msgs::msg::Header& header,
+                                  const cv::Mat& img_in,
+                                  const cv::Mat& color_in,
+                                  const CmapFunc& colormap,
+                                  const DisplayConfig& config) {
+  if (color_in.empty()) {
+    return makeImage(header, img_in, colormap, config);
+  }
+
+  cv::Mat img(img_in.rows, img_in.cols, CV_8UC3);
+  for (int r = 0; r < img_in.rows; ++r) {
+    for (int c = 0; c < img_in.cols; ++c) {
+      auto pixel = img.ptr<uint8_t>(r, c);
+      const auto c_in = colormap(img_in, r, c);
+      const auto c_cv = color_in.at<cv::Vec3b>(r, c);
+      const auto color =
+          c_in.blend(spark_dsg::Color(c_cv[0], c_cv[1], c_cv[2]), config.overlay_alpha);
+
+      *pixel = color.r;
+      *(pixel + 1) = color.g;
+      *(pixel + 2) = color.b;
+    }
+  }
+
+  return convertImage(header, img, config);
 }
 
-Image::SharedPtr makeRangeImage(const std_msgs::msg::Header& header,
-                                const InputData& sensor_data,
-                                const DisplayConfig& config) {
-  cv_bridge::CvImagePtr msg(new cv_bridge::CvImage());
-  msg->header = header;
-  msg->encoding = "32FC1";
-  msg->image = showImage(sensor_data.range_image, config);
-  return msg->toImageMsg();
+Image::SharedPtr makeDistImage(const std_msgs::msg::Header& header,
+                               const cv::Mat& distances,
+                               const DisplayConfig& config) {
+  double min_v = 0.0;
+  double max_v = std::numeric_limits<double>::infinity();
+  cv::minMaxIdx(distances, &min_v, &max_v);
+  const float v_min = config.min_distance >= 0.0f ? config.min_distance : min_v;
+  const float v_max = config.max_distance >= 0.0f ? config.max_distance : max_v;
+  const visualizer::RangeColormap cmap(config.distance_colormap);
+  return makeImage(
+      header,
+      distances,
+      [&](const cv::Mat& img, int r, int c) {
+        return cmap(img.at<float>(r, c), v_min, v_max);
+      },
+      config);
 }
 
 // TODO(nathan) pcl_ros would avoid this but it causes compile issues (and would also
