Ray casting-based depth image renderer

examples/python/raycast/depth_exctraction.py

@@ -0,0 +1,57 @@
+"""
+Depth map extraction from HashedWaveletOctree map at given camera pose and intrinsics
+"""
+
+import time
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+import pywavemap as wm
+
+
+def save_depth_as_png(depth_map: np.ndarray, output_path: Path):
+    depth_min = np.min(depth_map)
+    depth_max = np.max(depth_map)
+
+    # Avoid division by zero in case all values are the same
+    if depth_max - depth_min > 0:
+        depth_map_normalized = (depth_map - depth_min) / (depth_max -
+                                                          depth_min)
+    else:
+        depth_map_normalized = np.zeros_like(depth_map)
+
+    # Convert floats (meters) to uint8 and save to png
+    depth_map_8bit = (depth_map_normalized * 255).astype(np.uint8)
+    image = Image.fromarray(depth_map_8bit)
+    image.save(output_path)
+
+
+if __name__ == "__main__":
+    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
+
+    # Load map from file
+    your_map = wm.Map.load(map_path)
+
+    # 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)
+    t2 = time.perf_counter()
+    print(f"Depth your_map of size {camera_cfg.width}x{camera_cfg.height} "
+          f"created in {t2 - t1:.02f} seconds")
+    save_depth_as_png(depth, out_path)

library/cpp/include/wavemap/core/utils/math/int_math.h

@@ -111,6 +111,16 @@ constexpr T div_round_up(T numerator, T denominator) {
   return numerator / denominator + (numerator % denominator != 0);
 }
 
+template <int dim>
+Eigen::Matrix<int, dim, 1> div_round_up(Eigen::Matrix<int, dim, 1> vector,
+                                        int denominator) {
+  DCHECK_GE(denominator, 0);
+  for (int dim_idx = 0; dim_idx < dim; ++dim_idx) {
+    vector[dim_idx] = div_round_up(vector[dim_idx], denominator);
+  }
+  return vector;
+}
+
 template <typename T>
 constexpr T factorial(T x) {
   T result = 1;

library/python/CMakeLists.txt

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

library/python/include/pywavemap/raycast.h

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

library/python/src/pywavemap.cc

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

library/python/src/pywavemap/__init__.py

@@ -9,6 +9,7 @@
                                   HashedChunkedWaveletOctree,
                                   InterpolationMode)
 from ._pywavemap_bindings import Pipeline
+from ._pywavemap_bindings import raycast, PinholeCameraProjectorConfig, get_depth
 
 # Binding submodules
 from ._pywavemap_bindings import logging, param, convert

library/python/src/raycast.cc

@@ -0,0 +1,118 @@
+#include "pywavemap/raycast.h"
+
+#include <nanobind/eigen/dense.h>  // to use eigen2numpy seamlessly
+#include <wavemap/core/common.h>
+#include <wavemap/core/data_structure/image.h>
+#include <wavemap/core/integrator/projection_model/pinhole_camera_projector.h>
+#include <wavemap/core/map/hashed_wavelet_octree.h>
+#include <wavemap/core/utils/iterate/grid_iterator.h>
+#include <wavemap/core/utils/query/query_accelerator.h>
+#include <wavemap/core/utils/thread_pool.h>
+
+#include "wavemap/core/utils/iterate/ray_iterator.h"
+
+using namespace nb::literals;  // NOLINT
+
+namespace wavemap {
+FloatingPoint raycast(const HashedWaveletOctree& map,
+                      const Point3D& start_point, const Point3D& end_point,
+                      FloatingPoint log_odds_threshold) {
+  const FloatingPoint mcw = map.getMinCellWidth();
+  const Ray ray(start_point, end_point, mcw);
+  for (const Index3D& ray_voxel_index : ray) {
+    if (log_odds_threshold < map.getCellValue(ray_voxel_index)) {
+      const Point3D voxel_center =
+          convert::indexToCenterPoint(ray_voxel_index, mcw);
+      return (voxel_center - start_point).norm();
+    }
+  }
+  return (end_point - start_point).norm();
+}
+
+FloatingPoint raycast_fast(
+    QueryAccelerator<HashedWaveletOctree>& query_accelerator,
+    const Point3D& start_point, const Point3D& end_point,
+    FloatingPoint log_odds_threshold) {
+  const FloatingPoint min_cell_width = query_accelerator.getMinCellWidth();
+  const Ray ray(start_point, end_point, min_cell_width);
+  for (const Index3D& ray_voxel_index : ray) {
+    if (log_odds_threshold < query_accelerator.getCellValue(ray_voxel_index)) {
+      const Point3D voxel_center =
+          convert::indexToCenterPoint(ray_voxel_index, min_cell_width);
+      return (voxel_center - start_point).norm();
+    }
+  }
+  return (end_point - start_point).norm();
+}
+
+void add_raycast_bindings(nb::module_& m) {
+  nb::class_<PinholeCameraProjectorConfig>(
+      m, "PinholeCameraProjectorConfig", "Describes pinhole camera intrinsics")
+      .def(nb::init<FloatingPoint, FloatingPoint, FloatingPoint, FloatingPoint,
+                    IndexElement, IndexElement>(),
+           "fx"_a, "fy"_a, "cx"_a, "cy"_a, "height"_a, "width"_a)
+      .def_rw("width", &PinholeCameraProjectorConfig::width)
+      .def_rw("height", &PinholeCameraProjectorConfig::height)
+      .def_rw("fx", &PinholeCameraProjectorConfig::fx)
+      .def_rw("fy", &PinholeCameraProjectorConfig::fy)
+      .def_rw("cx", &PinholeCameraProjectorConfig::cx)
+      .def_rw("cy", &PinholeCameraProjectorConfig::cy)
+      .def("__repr__", [](const PinholeCameraProjectorConfig& self) {
+        return nb::str(
+                   "PinholeCameraProjectorConfig(width={}, height={}, fx={}, "
+                   "fy={}, cx={}, cy={})")
+            .format(self.width, self.height, self.fx, self.fy, self.cx,
+                    self.cy);
+      });
+
+  m.def("raycast", &raycast,
+        "Raycast and get first point with occopancy higher than threshold");
+
+  m.def("raycast_fast",
+        &raycast_fast,  // TODO: unusable without QueryAccelerator binding
+        "Raycast and get first point with occopancy higher than threshold "
+        "using QueryAccelerator for efficiency");
+
+  m.def(
+      "get_depth",
+      [](const HashedWaveletOctree& map, const Transformation3D& pose,
+         const PinholeCameraProjectorConfig& cam_cfg,
+         FloatingPoint log_odds_threshold, FloatingPoint max_range) {
+        const PinholeCameraProjector projection_model(cam_cfg);
+        Image depth_image(projection_model.getDimensions());
+        const Point3D& W_start_point = pose.getPosition();
+
+        ThreadPool thread_pool;
+        constexpr int kPatchWidth = 64;
+        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([&map, &projection_model, &pose, &W_start_point,
+                                &depth_image, patch_index, kPatchWidth,
+                                max_range, log_odds_threshold]() {
+            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<2>(patch_min, patch_max)) {
+              FloatingPoint& depth_pixel = depth_image.at(index);
+              const Vector2D image_xy = projection_model.indexToImage(index);
+              const Point3D C_end_point =
+                  projection_model.sensorToCartesian({image_xy, max_range});
+              const Point3D W_end_point = pose * C_end_point;
+              depth_pixel = raycast_fast(query_accelerator, W_start_point,
+                                         W_end_point, log_odds_threshold);
+            }
+          });
+        }
+        thread_pool.wait_all();
+
+        return depth_image.getData().transpose().eval();
+      },
+      "Extract depth from octree map at using given camera pose and "
+      "intrinsics");
+}
+}  // namespace wavemap