Support multi layer images (#79)

### What change is being made

Please provide a detailed description of WHAT change is being made such
that the reader is able to
understand the change holistically.

### Why this change is being made

Please provide the rationale behind the change and additional context
like links to documents,
related work or issues.

### Tested

Please provide a description how this change was tested, e.g. unit
tests, hardware tests or commands
you run.

GitOrigin-RevId: 07952e29e2a1b88ea2c3659fb1479bc6d87ef457

Author: awoll-bdai

control/components.cpp

@@ -352,57 +352,78 @@ bool HeightScanInput::read(OnnxRuntime& runtime, RobotStateInterface& state, Com
       LOG_STREAM(ERROR, fmt::format("Failed to get input buffer {}.{}", key_, layer_name));
       return false;
     }
-    copyToBuffer(maybe_scan.value()->layers.at(layer_name), maybe_buffer.value());
+    copyToBuffer(maybe_scan.value()->float_layers.at(layer_name), maybe_buffer.value());
   }
   return true;
 }
 
-// Implementation of RangeImageInput methods
-RangeImageInput::RangeImageInput(const std::string& key,
-                                 const metadata::RangeImageMetadata& metadata)
-    : key_(key), metadata_(metadata) {}
+// Implementation of SphericalImageInput methods
+SphericalImageInput::SphericalImageInput(const std::string& key, const std::string& sensor_name,
+                                         const std::unordered_set<std::string>& channel_names,
+                                         const metadata::SphericalImageMetadata& metadata)
+    : key_(key), sensor_name_(sensor_name), channel_names_(channel_names), metadata_(metadata) {}
 
-bool RangeImageInput::init(RobotStateInterface& state, CommandInterface&) {
+bool SphericalImageInput::init(RobotStateInterface& state, CommandInterface&) {
   SphericalImageConfig config;
   config.v_res = static_cast<int>(metadata_.v_res);
   config.h_res = static_cast<int>(metadata_.h_res);
   config.v_fov_min_deg = metadata_.v_fov_min_deg;
   config.v_fov_max_deg = metadata_.v_fov_max_deg;
   config.unobserved_value = metadata_.unobserved_value;
-  return state.initRangeImage(config);
+  config.channel_names = channel_names_;
+  return state.initSphericalImage(sensor_name_, config);
 }
 
-bool RangeImageInput::read(OnnxRuntime& runtime, RobotStateInterface& state, CommandInterface&) {
-  auto maybe_buffer = runtime.inputBuffer<float>(key_);
-  if (!maybe_buffer.has_value()) return false;
-  auto maybe_image = state.rangeImage();
-  if (!maybe_image.has_value()) return false;
-  copyToBuffer(maybe_image.value(), maybe_buffer.value());
+bool SphericalImageInput::read(OnnxRuntime& runtime, RobotStateInterface& state,
+                               CommandInterface&) {
+  auto maybe_image = state.sphericalImage(sensor_name_, channel_names_);
+  if (!maybe_image.has_value()) {
+    LOG_STREAM(ERROR, "Failed to get spherical image data for SphericalImageInput");
+    return false;
+  }
+  for (const auto& channel_name : channel_names_) {
+    auto maybe_buffer = runtime.inputBuffer<float>(fmt::format("{}.{}", key_, channel_name));
+    if (!maybe_buffer.has_value()) {
+      LOG_STREAM(ERROR, fmt::format("Failed to get input buffer {}.{}", key_, channel_name));
+      return false;
+    }
+    copyToBuffer(maybe_image.value()->float_channels.at(channel_name), maybe_buffer.value());
+  }
   return true;
 }
 
-// Implementation of DepthImageInput methods
-DepthImageInput::DepthImageInput(const std::string& key,
-                                 const metadata::DepthImageMetadata& metadata)
-    : key_(key), metadata_(metadata) {}
+// Implementation of PinholeImageInput methods
+PinholeImageInput::PinholeImageInput(const std::string& key, const std::string& sensor_name,
+                                     const std::unordered_set<std::string>& channel_names,
+                                     const metadata::PinholeImageMetadata& metadata)
+    : key_(key), sensor_name_(sensor_name), channel_names_(channel_names), metadata_(metadata) {}
 
-bool DepthImageInput::init(RobotStateInterface& state, CommandInterface&) {
+bool PinholeImageInput::init(RobotStateInterface& state, CommandInterface&) {
   PinholeImageConfig config;
   config.width = metadata_.width;
   config.height = metadata_.height;
   config.fx = metadata_.fx;
   config.fy = metadata_.fy;
   config.cx = metadata_.cx;
   config.cy = metadata_.cy;
-  return state.initDepthImage(config);
+  config.channel_names = channel_names_;
+  return state.initPinholeImage(sensor_name_, config);
 }
 
-bool DepthImageInput::read(OnnxRuntime& runtime, RobotStateInterface& state, CommandInterface&) {
-  auto maybe_buffer = runtime.inputBuffer<float>(key_);
-  if (!maybe_buffer.has_value()) return false;
-  auto maybe_image = state.depthImage();
-  if (!maybe_image.has_value()) return false;
-  copyToBuffer(maybe_image.value(), maybe_buffer.value());
+bool PinholeImageInput::read(OnnxRuntime& runtime, RobotStateInterface& state, CommandInterface&) {
+  auto maybe_image = state.pinholeImage(sensor_name_, channel_names_);
+  if (!maybe_image.has_value()) {
+    LOG_STREAM(ERROR, "Failed to get pinhole image data for PinholeImageInput");
+    return false;
+  }
+  for (const auto& channel_name : channel_names_) {
+    auto maybe_buffer = runtime.inputBuffer<float>(fmt::format("{}.{}", key_, channel_name));
+    if (!maybe_buffer.has_value()) {
+      LOG_STREAM(ERROR, fmt::format("Failed to get input buffer {}.{}", key_, channel_name));
+      return false;
+    }
+    copyToBuffer(maybe_image.value()->float_channels.at(channel_name), maybe_buffer.value());
+  }
   return true;
 }
 

control/components.hpp

@@ -299,52 +299,67 @@ class HeightScanInput : public Input {
 };
 
 /**
- * @brief Input component that reads range/depth image from LiDAR sensor.
+ * @brief Input component that reads spherical image data with multiple channels.
  *
- * Reads range image data (distances) from a LiDAR sensor and copies it to the
- * ONNX input buffer. Configuration includes resolution, field of view, and
- * sentinel value for unobserved points.
+ * Reads spherical image data from a named sensor with configurable channels (e.g., range,
+ * risk). Each channel is written to a separate ONNX input buffer with the
+ * naming pattern "key.channel_name".
  */
-class RangeImageInput : public Input {
+class SphericalImageInput : public Input {
  public:
   /**
-   * @brief Construct a range image input component.
+   * @brief Construct a spherical image input component.
    *
-   * @param key ONNX input tensor name (e.g., "sensors.lidar.range").
-   * @param metadata Range image metadata containing resolution, FOV, and sentinel value.
+   * @param key ONNX input tensor base name (e.g., "sensor.spherical_image.lidar1").
+   * @param sensor_name Name of the spherical image sensor to read from.
+   * @param channel_names Set of channel names to include (each creates a buffer
+   * "key.channel_name").
+   * @param metadata Spherical image metadata containing resolution, FOV, and sentinel value.
    */
-  RangeImageInput(const std::string& key, const metadata::RangeImageMetadata& metadata);
+  SphericalImageInput(const std::string& key, const std::string& sensor_name,
+                      const std::unordered_set<std::string>& channel_names,
+                      const metadata::SphericalImageMetadata& metadata);
 
   bool init(RobotStateInterface& state, CommandInterface& command) override;
   bool read(OnnxRuntime& runtime, RobotStateInterface& state, CommandInterface& command) override;
 
  private:
-  std::string key_;                        ///< ONNX input tensor name.
-  metadata::RangeImageMetadata metadata_;  ///< Range image configuration.
+  std::string key_;                                ///< ONNX input tensor base name.
+  std::string sensor_name_;                        ///< Spherical image sensor name.
+  std::unordered_set<std::string> channel_names_;  ///< Channel names to read.
+  metadata::SphericalImageMetadata metadata_;      ///< Spherical image configuration.
 };
 
 /**
- * @brief Input component that reads depth image from camera sensor.
+ * @brief Input component that reads pinhole camera image data with multiple channels.
  *
- * Reads depth image data from a camera sensor and copies it to the ONNX input
- * buffer. Configuration includes image dimensions and camera intrinsic parameters.
+ * Reads pinhole image data from a named sensor with configurable channels (e.g., depth,
+ * risk). Each channel is written to a separate ONNX input buffer with the
+ * naming pattern "key.channel_name".
  */
-class DepthImageInput : public Input {
+class PinholeImageInput : public Input {
  public:
   /**
-   * @brief Construct a depth image input component.
+   * @brief Construct a pinhole image input component.
    *
-   * @param key ONNX input tensor name (e.g., "sensors.camera.depth").
-   * @param metadata Depth image metadata containing width, height, and camera intrinsics.
+   * @param key ONNX input tensor base name (e.g., "sensor.pinhole_image.cam1").
+   * @param sensor_name Name of the pinhole image sensor to read from.
+   * @param channel_names Set of channel names to include (each creates a buffer
+   * "key.channel_name").
+   * @param metadata Pinhole image metadata containing width, height, and camera intrinsics.
    */
-  DepthImageInput(const std::string& key, const metadata::DepthImageMetadata& metadata);
+  PinholeImageInput(const std::string& key, const std::string& sensor_name,
+                    const std::unordered_set<std::string>& channel_names,
+                    const metadata::PinholeImageMetadata& metadata);
 
   bool init(RobotStateInterface& state, CommandInterface& command) override;
   bool read(OnnxRuntime& runtime, RobotStateInterface& state, CommandInterface& command) override;
 
  private:
-  std::string key_;                        ///< ONNX input tensor name.
-  metadata::DepthImageMetadata metadata_;  ///< Depth image configuration.
+  std::string key_;                                ///< ONNX input tensor base name.
+  std::string sensor_name_;                        ///< Pinhole image sensor name.
+  std::unordered_set<std::string> channel_names_;  ///< Channel names to read.
+  metadata::PinholeImageMetadata metadata_;        ///< Pinhole image configuration.
 };
 
 /**

control/controller.cpp

@@ -23,8 +23,6 @@ OnnxRLController::OnnxRLController(RobotStateInterface& state, CommandInterface&
   context_.registerMatcher(std::make_unique<SE2VelocityMatcher>());
   context_.registerMatcher(std::make_unique<IMUAngularVelocityMatcher>());
   context_.registerMatcher(std::make_unique<IMUOrientationMatcher>());
-  context_.registerMatcher(std::make_unique<RangeImageMatcher>());
-  context_.registerMatcher(std::make_unique<DepthImageMatcher>());
   context_.registerMatcher(std::make_unique<BodyPositionMatcher>());
   context_.registerMatcher(std::make_unique<BodyOrientationMatcher>());
   context_.registerMatcher(std::make_unique<CommandSE3PoseMatcher>());
@@ -36,6 +34,8 @@ OnnxRLController::OnnxRLController(RobotStateInterface& state, CommandInterface&
   context_.registerGroupMatcher(std::make_unique<JointMatcher>());
   context_.registerGroupMatcher(std::make_unique<JointTargetMatcher>());
   context_.registerGroupMatcher(std::make_unique<HeightScanMatcher>());
+  context_.registerGroupMatcher(std::make_unique<SphericalImageMatcher>());
+  context_.registerGroupMatcher(std::make_unique<PinholeImageMatcher>());
   context_.registerGroupMatcher(std::make_unique<MemoryMatcher>());
 }
 

control/interfaces.hpp

@@ -5,8 +5,10 @@
 #include <Eigen/Geometry>
 
 #include <optional>
+#include <span>
 #include <string>
 #include <string_view>
+#include <unordered_map>
 #include <vector>
 
 /**
@@ -101,12 +103,28 @@ struct SE2VelocityRanges {
  */
 struct HeightScan {
   /**
-   * @brief Map of layer names to their data vectors.
+   * @brief Map of layer names to their data spans.
    *
-   * Each entry maps a layer name (e.g., "height", "color") to a vector of
-   * double values representing the flattened grid data for that layer.
+   * Each entry maps a layer name (e.g., "height", "color") to a span of
+   * float values representing the flattened grid data for that layer.
    */
-  std::unordered_map<std::string, std::vector<double>> layers;
+  std::unordered_map<std::string, std::span<const float>> float_layers;
+};
+
+/**
+ * @brief A flattened multi-channel image.
+ *
+ * Each channel represents a different data type (e.g., depth, range, risk).
+ * All channels share the same projection and must have the same length.
+ */
+struct MultiChannelImage {
+  /**
+   * @brief Map of channel names to their data spans.
+   *
+   * Each entry maps a channel name to a span of float values representing the flattened image data
+   * for that channel.
+   */
+  std::unordered_map<std::string, std::span<const float>> float_channels;
 };
 
 }  // namespace exploy::control