Add map edit module to Python API and expose crop_to_sphere methods

ethz-asl · Dec 12, 2024 · a815bc0 · a815bc0
1 parent f37f7f9
commit a815bc0
Show file tree

Hide file tree

Showing 7 changed files with 83 additions and 12 deletions.
diff --git a/examples/cpp/edit/transform_map.cc b/examples/cpp/edit/transform_map.cc
@@ -5,16 +5,10 @@
 #include <wavemap/io/file_conversions.h>
 
 int main(int, char**) {
-  // Settings
+  // Load the map
   std::filesystem::path input_map_path =
       "/home/victor/data/wavemaps/newer_college_cloister_10cm.wvmp";
-  std::filesystem::path output_map_path =
-      "/home/victor/data/wavemaps/newer_college_cloister_10cm_tranformed.wvmp";
-
-  // Create a smart pointer that will own the loaded map
   wavemap::MapBase::Ptr map_base;
-
-  // Load the map
   bool success = wavemap::io::fileToMap(input_map_path, map_base);
   CHECK(success);
 
@@ -31,6 +25,8 @@ int main(int, char**) {
                                  std::make_shared<wavemap::ThreadPool>());
 
   // Save the map
+  std::filesystem::path output_map_path =
+      "/home/victor/data/wavemaps/newer_college_cloister_10cm_tranformed.wvmp";
   success &= wavemap::io::mapToFile(*map, output_map_path);
   CHECK(success);
 }
diff --git a/examples/python/edit/crop_map.py b/examples/python/edit/crop_map.py
@@ -0,0 +1,17 @@
+import os
+import numpy as np
+import pywavemap as wave
+
+# Load the map
+user_home = os.path.expanduser('~')
+input_map_path = os.path.join(user_home, "your_map.wvmp")
+your_map = wave.Map.load(input_map_path)
+
+# Crop the map
+center = np.array([-2.2, -1.4, 0.0])
+radius = 3.0
+wave.edit.crop_to_sphere(your_map, center, radius)
+
+# Save the map
+output_map_path = os.path.join(user_home, "your_map_cropped.wvmp")
+your_map.store(output_map_path)
diff --git a/library/python/CMakeLists.txt b/library/python/CMakeLists.txt
@@ -58,6 +58,7 @@ find_package(nanobind CONFIG REQUIRED)
 nanobind_add_module(_pywavemap_bindings STABLE_ABI
     src/pywavemap.cc
     src/convert.cc
+    src/edit.cc
     src/indices.cc
     src/logging.cc
     src/maps.cc
@@ -81,10 +82,6 @@ nanobind_add_stub(pywavemap_stub INSTALL_TIME
     OUTPUT "pywavemap/__init__.pyi"
     MARKER_FILE "pywavemap/py.typed"
     PYTHON_PATH "pywavemap")
-nanobind_add_stub(pywavemap_convert_stub INSTALL_TIME
-    MODULE _pywavemap_bindings.convert
-    OUTPUT "pywavemap/convert.pyi"
-    PYTHON_PATH "pywavemap")
 nanobind_add_stub(pywavemap_logging_stub INSTALL_TIME
     MODULE _pywavemap_bindings.logging
     OUTPUT "pywavemap/logging.pyi"
@@ -93,3 +90,11 @@ nanobind_add_stub(pywavemap_param_stub INSTALL_TIME
     MODULE _pywavemap_bindings.param
     OUTPUT "pywavemap/param.pyi"
     PYTHON_PATH "pywavemap")
+nanobind_add_stub(pywavemap_convert_stub INSTALL_TIME
+    MODULE _pywavemap_bindings.convert
+    OUTPUT "pywavemap/convert.pyi"
+    PYTHON_PATH "pywavemap")
+nanobind_add_stub(pywavemap_edit_stub INSTALL_TIME
+    MODULE _pywavemap_bindings.edit
+    OUTPUT "pywavemap/edit.pyi"
+    PYTHON_PATH "pywavemap")
diff --git a/library/python/include/pywavemap/edit.h b/library/python/include/pywavemap/edit.h
@@ -0,0 +1,12 @@
+#ifndef PYWAVEMAP_EDIT_H_
+#define PYWAVEMAP_EDIT_H_
+
+#include <nanobind/nanobind.h>
+
+namespace nb = nanobind;
+
+namespace wavemap {
+void add_edit_module(nb::module_& m_edit);
+}  // namespace wavemap
+
+#endif  // PYWAVEMAP_EDIT_H_
diff --git a/library/python/src/edit.cc b/library/python/src/edit.cc
@@ -0,0 +1,32 @@
+#include "pywavemap/edit.h"
+
+#include <memory>
+
+#include <nanobind/eigen/dense.h>
+#include <wavemap/core/map/hashed_chunked_wavelet_octree.h>
+#include <wavemap/core/map/hashed_wavelet_octree.h>
+#include <wavemap/core/utils/edit/crop.h>
+
+using namespace nb::literals;  // NOLINT
+
+namespace wavemap {
+void add_edit_module(nb::module_& m_edit) {
+  // Map cropping methods
+  m_edit.def(
+      "crop_to_sphere",
+      [](HashedWaveletOctree& map, const Point3D& t_W_center,
+         FloatingPoint radius, IndexElement termination_height) {
+        edit::crop_to_sphere(map, t_W_center, radius, termination_height,
+                             std::make_shared<ThreadPool>());
+      },
+      "map"_a, "center_point"_a, "radius"_a, "termination_height"_a = 0);
+  m_edit.def(
+      "crop_to_sphere",
+      [](HashedChunkedWaveletOctree& map, const Point3D& t_W_center,
+         FloatingPoint radius, IndexElement termination_height) {
+        edit::crop_to_sphere(map, t_W_center, radius, termination_height,
+                             std::make_shared<ThreadPool>());
+      },
+      "map"_a, "center_point"_a, "radius"_a, "termination_height"_a = 0);
+}
+}  // namespace wavemap
diff --git a/library/python/src/pywavemap.cc b/library/python/src/pywavemap.cc
@@ -1,6 +1,7 @@
 #include <nanobind/nanobind.h>
 
 #include "pywavemap/convert.h"
+#include "pywavemap/edit.h"
 #include "pywavemap/indices.h"
 #include "pywavemap/logging.h"
 #include "pywavemap/maps.h"
@@ -42,6 +43,14 @@ NB_MODULE(_pywavemap_bindings, m) {
       "Submodule with common conversion functions for wavemap index types.");
   add_convert_module(m_convert);
 
+  // Bindings for map editing tools
+  nb::module_ m_edit =
+      m.def_submodule("edit",
+                      "edit\n"
+                      "=======\n"
+                      "Submodule with tools to edit wavemap maps.");
+  add_edit_module(m_edit);
+
   // Bindings for index types
   add_index_bindings(m);
 

diff --git a/library/python/src/pywavemap/__init__.py b/library/python/src/pywavemap/__init__.py
@@ -11,4 +11,4 @@
 from ._pywavemap_bindings import Pipeline
 
 # Binding submodules
-from ._pywavemap_bindings import logging, param, convert
+from ._pywavemap_bindings import logging, param, convert, edit