Move code to C++ library and add support for other projection models

examples/python/raycast/depth_exctraction.py → examples/python/render/raycast_depth_image.py

@@ -1,5 +1,5 @@
 """
-Depth map extraction from HashedWaveletOctree map at given camera pose and intrinsics
+Extract depth maps from wavemap maps at a given sensor pose and intrinsics
 """
 
 import time
@@ -31,27 +31,43 @@ def save_depth_as_png(depth_map: np.ndarray, output_path: Path):
     map_path = Path.home() \
                / "data/panoptic_mapping/flat_dataset/run2/your_map.wvmp"
     out_path = Path(__file__).parent / "depth.png"
-    camera_cfg = wm.PinholeCameraProjectorConfig(
-        width=1280,
-        height=720,
-        fx=526.21539307,
-        fy=526.21539307,
-        cx=642.309021,
-        cy=368.69949341,
-    )  # Note: these are intrinsics for Zed 2i
+
+    # Configure the virtual sensor's projection model
+    # NOTE: These intrinsics match the Zed 2i depth camera.
+    projection_model = wm.Projector.create({
+        "type": "pinhole_camera_projector",
+        "width": 1280,
+        "height": 720,
+        "fx": 526.21539307,
+        "fy": 526.21539307,
+        "cx": 642.309021,
+        "cy": 368.69949341
+    })
 
     # Load map from file
     your_map = wm.Map.load(map_path)
 
+    # Create the depth image renderer
+    log_odds_occupancy_threshold = 0.1
+    min_range = 0.1
+    max_range = 6.0
+    default_depth_value = -1.0
+    renderer = wm.RaycastingRenderer(your_map, projection_model,
+                                     log_odds_occupancy_threshold, min_range,
+                                     max_range, default_depth_value)
+
     # Create pose
     rotation = wm.Rotation(np.eye(3))
     translation = np.zeros(3)
     pose = wm.Pose(rotation, translation)
 
     # Extract depth
     t1 = time.perf_counter()
-    depth = wm.get_depth(your_map, pose, camera_cfg, 0.1, 10)
+    depth_image = renderer.render(pose)
     t2 = time.perf_counter()
-    print(f"Depth your_map of size {camera_cfg.width}x{camera_cfg.height} "
+    print(f"Depth image of size {depth_image.width}x{depth_image.height} "
           f"created in {t2 - t1:.02f} seconds")
-    save_depth_as_png(depth, out_path)
+    save_depth_as_png(depth_image.data, out_path)
+
+    # Avoids leak warnings on old Python versions with lazy garbage collectors
+    del renderer, depth_image

library/cpp/include/wavemap/core/utils/render/impl/raycast_inl.h

@@ -0,0 +1,34 @@
+#ifndef WAVEMAP_CORE_UTILS_RENDER_IMPL_RAYCAST_INL_H_
+#define WAVEMAP_CORE_UTILS_RENDER_IMPL_RAYCAST_INL_H_
+
+namespace wavemap::raycast {
+template <typename MapT>
+std::optional<Index3D> first_collision_index(
+    MapT& map, const Point3D& W_start_point, const Point3D& W_end_point,
+    FloatingPoint log_odds_occupancy_threshold) {
+  const Ray ray(W_start_point, W_end_point, map.getMinCellWidth());
+  for (const Index3D& ray_voxel_index : ray) {
+    if (log_odds_occupancy_threshold < map.getCellValue(ray_voxel_index)) {
+      return ray_voxel_index;
+    }
+  }
+  return std::nullopt;
+}
+
+template <typename MapT>
+std::optional<FloatingPoint> first_collision_distance(
+    MapT& map, const Point3D& W_start_point, const Point3D& W_end_point,
+    FloatingPoint log_odds_occupancy_threshold) {
+  const auto colliding_index = first_collision_index<MapT>(
+      map, W_start_point, W_end_point, log_odds_occupancy_threshold);
+  if (colliding_index) {
+    const FloatingPoint min_cell_width = map.getMinCellWidth();
+    const Point3D voxel_center =
+        convert::indexToCenterPoint(colliding_index.value(), min_cell_width);
+    return (voxel_center - W_start_point).norm();
+  }
+  return std::nullopt;
+}
+}  // namespace wavemap::raycast
+
+#endif  // WAVEMAP_CORE_UTILS_RENDER_IMPL_RAYCAST_INL_H_

library/cpp/include/wavemap/core/utils/render/raycast.h

@@ -0,0 +1,25 @@
+#ifndef WAVEMAP_CORE_UTILS_RENDER_RAYCAST_H_
+#define WAVEMAP_CORE_UTILS_RENDER_RAYCAST_H_
+
+#include "wavemap/core/common.h"
+#include "wavemap/core/utils/iterate/ray_iterator.h"
+
+namespace wavemap::raycast {
+// NOTE: The map can either be a raw map, or a QueryAccelerator instance. We
+//       strongly recommend passing a QueryAccelerator to get good performance.
+template <typename MapT>
+std::optional<Index3D> first_collision_index(
+    MapT& map, const Point3D& W_start_point, const Point3D& W_end_point,
+    FloatingPoint log_odds_occupancy_threshold);
+
+// NOTE: The map can either be a raw map, or a QueryAccelerator instance. We
+//       strongly recommend passing a QueryAccelerator to get good performance.
+template <typename MapT>
+std::optional<FloatingPoint> first_collision_distance(
+    MapT& map, const Point3D& W_start_point, const Point3D& W_end_point,
+    FloatingPoint log_odds_occupancy_threshold);
+}  // namespace wavemap::raycast
+
+#include "wavemap/core/utils/render/impl/raycast_inl.h"
+
+#endif  // WAVEMAP_CORE_UTILS_RENDER_RAYCAST_H_

library/cpp/include/wavemap/core/utils/render/raycasting_renderer.h

@@ -0,0 +1,43 @@
+#ifndef WAVEMAP_CORE_UTILS_RENDER_RAYCASTING_RENDERER_H_
+#define WAVEMAP_CORE_UTILS_RENDER_RAYCASTING_RENDERER_H_
+
+#include <memory>
+
+#include "wavemap/core/data_structure/image.h"
+#include "wavemap/core/integrator/projection_model/projector_base.h"
+#include "wavemap/core/map/map_base.h"
+#include "wavemap/core/utils/thread_pool.h"
+
+namespace wavemap {
+class RaycastingRenderer {
+ public:
+  RaycastingRenderer(MapBase::ConstPtr occupancy_map,
+                     ProjectorBase::ConstPtr projection_model,
+                     FloatingPoint log_odds_occupancy_threshold = 1e-3f,
+                     FloatingPoint min_range = 0.f,
+                     FloatingPoint max_range = 10.f,
+                     FloatingPoint default_depth_value = -1.f,
+                     std::shared_ptr<ThreadPool> thread_pool = nullptr);
+
+  const Image<>& render(const Transformation3D& T_W_C);
+
+ private:
+  static constexpr int kPatchWidth = 64;
+
+  const MapBase::ConstPtr map_;
+  const std::shared_ptr<ThreadPool> thread_pool_;
+
+  const ProjectorBase::ConstPtr projection_model_;
+  const FloatingPoint min_range_;
+  const FloatingPoint max_range_;
+  const FloatingPoint log_odds_occupancy_threshold_;
+
+  Image<FloatingPoint> depth_image_;
+
+  template <typename MapT>
+  void renderPatch(const MapT& map, const Transformation3D& T_W_C,
+                   const Index2D& patch_index);
+};
+}  // namespace wavemap
+
+#endif  // WAVEMAP_CORE_UTILS_RENDER_RAYCASTING_RENDERER_H_

library/cpp/src/core/CMakeLists.txt

@@ -46,6 +46,7 @@ target_sources(wavemap_core PRIVATE
     utils/query/classified_map.cc
     utils/query/query_accelerator.cc
     utils/query/point_sampler.cc
+    utils/render/raycasting_renderer.cc
     utils/sdf/full_euclidean_sdf_generator.cc
     utils/sdf/quasi_euclidean_sdf_generator.cc
     utils/time/stopwatch.cc

library/cpp/src/core/utils/render/raycasting_renderer.cc

@@ -0,0 +1,78 @@
+#include "wavemap/core/utils/render/raycasting_renderer.h"
+
+#include <memory>
+#include <utility>
+
+#include "wavemap/core/utils/iterate/grid_iterator.h"
+#include "wavemap/core/utils/query/query_accelerator.h"
+#include "wavemap/core/utils/render/raycast.h"
+
+namespace wavemap {
+RaycastingRenderer::RaycastingRenderer(
+    MapBase::ConstPtr occupancy_map, ProjectorBase::ConstPtr projection_model,
+    FloatingPoint log_odds_occupancy_threshold, FloatingPoint min_range,
+    FloatingPoint max_range, FloatingPoint default_depth_value,
+    std::shared_ptr<ThreadPool> thread_pool)
+    : map_(CHECK_NOTNULL(occupancy_map)),
+      thread_pool_(thread_pool ? std::move(thread_pool)
+                               : std::make_shared<ThreadPool>()),
+      projection_model_(CHECK_NOTNULL(projection_model)),
+      min_range_(min_range),
+      max_range_(max_range),
+      log_odds_occupancy_threshold_(log_odds_occupancy_threshold),
+      depth_image_(projection_model_->getDimensions(), default_depth_value) {}
+
+template <typename MapT>
+void RaycastingRenderer::renderPatch(const MapT& map,
+                                     const Transformation3D& T_W_C,
+                                     const Index2D& patch_index) {
+  QueryAccelerator query_accelerator(map);
+
+  const Index2D patch_min = patch_index * kPatchWidth;
+  const Index2D patch_max = (patch_min.array() + kPatchWidth)
+                                .min(depth_image_.getDimensions().array()) -
+                            1;
+
+  for (const Index2D& index : Grid(patch_min, patch_max)) {
+    FloatingPoint& depth_pixel = depth_image_.at(index);
+    const Vector2D image_xy = projection_model_->indexToImage(index);
+    const Point3D C_start_point =
+        projection_model_->sensorToCartesian({image_xy, min_range_});
+    const Point3D C_end_point =
+        projection_model_->sensorToCartesian({image_xy, max_range_});
+    const Point3D W_start_point = T_W_C * C_start_point;
+    const Point3D W_end_point = T_W_C * C_end_point;
+    const auto d_start_collision = raycast::first_collision_distance(
+        query_accelerator, W_start_point, W_end_point,
+        log_odds_occupancy_threshold_);
+    if (d_start_collision) {
+      depth_pixel = d_start_collision.value();
+    }
+  }
+}
+
+const Image<>& RaycastingRenderer::render(const Transformation3D& T_W_C) {
+  depth_image_.resetToInitialValue();
+
+  const Index2D num_patches =
+      int_math::div_round_up(depth_image_.getDimensions(), kPatchWidth);
+  for (const Index2D& patch_index :
+       Grid<2>(Index2D::Zero(), num_patches - Index2D::Ones())) {
+    thread_pool_->add_task([this, &T_W_C, patch_index]() {
+      if (const auto* map =
+              dynamic_cast<const HashedWaveletOctree*>(map_.get());
+          map) {
+        renderPatch(*map, T_W_C, patch_index);
+      }
+      if (const auto* map =
+              dynamic_cast<const HashedChunkedWaveletOctree*>(map_.get());
+          map) {
+        renderPatch(*map, T_W_C, patch_index);
+      }
+    });
+  }
+  thread_pool_->wait_all();
+
+  return depth_image_;
+}
+}  // namespace wavemap

library/python/CMakeLists.txt

@@ -64,7 +64,7 @@ nanobind_add_module(_pywavemap_bindings STABLE_ABI
     src/measurements.cc
     src/param.cc
     src/pipeline.cc
-    src/raycast.cc)
+    src/render.cc)
 set_wavemap_target_properties(_pywavemap_bindings)
 target_include_directories(_pywavemap_bindings PRIVATE include)
 target_link_libraries(_pywavemap_bindings PRIVATE

library/python/include/pywavemap/render.h

@@ -0,0 +1,12 @@
+#ifndef PYWAVEMAP_RENDER_H_
+#define PYWAVEMAP_RENDER_H_
+
+#include <nanobind/nanobind.h>
+
+namespace nb = nanobind;
+
+namespace wavemap {
+void add_render_bindings(nb::module_& m);
+}  // namespace wavemap
+
+#endif  // PYWAVEMAP_RENDER_H_

library/python/src/measurements.cc

@@ -1,9 +1,14 @@
 #include "pywavemap/measurements.h"
 
+#include <memory>
+
 #include <nanobind/eigen/dense.h>
+#include <nanobind/stl/shared_ptr.h>
 #include <wavemap/core/common.h>
 #include <wavemap/core/data_structure/image.h>
 #include <wavemap/core/data_structure/pointcloud.h>
+#include <wavemap/core/integrator/projection_model/projector_base.h>
+#include <wavemap/core/integrator/projection_model/projector_factory.h>
 
 using namespace nb::literals;  // NOLINT
 
@@ -33,9 +38,24 @@ void add_measurement_bindings(nb::module_& m) {
 
   // Images
   nb::class_<Image<>>(m, "Image", "A class to store depth images.")
-      .def(nb::init<Image<>::Data>(), "pixel_matrix"_a);
+      .def(nb::init<Image<>::Data>(), "pixel_matrix"_a)
+      .def_prop_ro("width", &Image<>::getNumRows,
+                   "The image's width in pixels.")
+      .def_prop_ro("height", &Image<>::getNumColumns,
+                   "The image's height in pixels.")
+      .def_prop_ro("data", nb::overload_cast<>(&Image<>::getData, nb::const_));
   nb::class_<PosedImage<>>(
       m, "PosedImage", "A class to store depth images with an associated pose.")
       .def(nb::init<Transformation3D, Image<>>(), "pose"_a, "image"_a);
+
+  // Projection models
+  nb::class_<ProjectorBase>(m, "Projector")
+      .def_static(
+          "create",
+          [](const param::Value& params) -> std::shared_ptr<ProjectorBase> {
+            return ProjectorFactory::create(params);
+          },
+          nb::sig("def create(parameters: dict) -> Projector"), "parameters"_a,
+          "Create a new projection model based on the given settings.");
 }
 }  // namespace wavemap

library/python/src/pywavemap.cc

@@ -7,7 +7,7 @@
 #include "pywavemap/measurements.h"
 #include "pywavemap/param.h"
 #include "pywavemap/pipeline.h"
-#include "pywavemap/raycast.h"
+#include "pywavemap/render.h"
 
 using namespace wavemap;  // NOLINT
 namespace nb = nanobind;
@@ -55,6 +55,6 @@ NB_MODULE(_pywavemap_bindings, m) {
   // Bindings for measurement integration and map update pipelines
   add_pipeline_bindings(m);
 
-  // Bindings for raycasting
-  add_raycast_bindings(m);
+  // Bindings for rendering
+  add_render_bindings(m);
 }

library/python/src/pywavemap/__init__.py

@@ -4,12 +4,12 @@
 # Binding types
 from ._pywavemap_bindings import OctreeIndex
 from ._pywavemap_bindings import (Rotation, Pose, Pointcloud, PosedPointcloud,
-                                  Image, PosedImage)
+                                  Image, PosedImage, Projector)
 from ._pywavemap_bindings import (Map, HashedWaveletOctree,
                                   HashedChunkedWaveletOctree,
                                   InterpolationMode)
 from ._pywavemap_bindings import Pipeline
-from ._pywavemap_bindings import raycast, PinholeCameraProjectorConfig, get_depth
+from ._pywavemap_bindings import RaycastingRenderer
 
 # Binding submodules
 from ._pywavemap_bindings import logging, param, convert