Fix/pyviz (#31)

* fix ids for open3d viz, add remove layer to bindings * fix layer id for names * cleanup nodes/edges when removing layer --------- Co-authored-by: Jared Strader <[email protected]>
MIT-SPARK · Feb 18, 2025 · c0b51a9 · c0b51a9
1 parent 83b1d9d
commit c0b51a9
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 17 deletions.
diff --git a/python/bindings/src/spark_dsg_bindings.cpp b/python/bindings/src/spark_dsg_bindings.cpp
@@ -139,7 +139,13 @@ PYBIND11_MODULE(_dsg_bindings, module) {
       .def_readonly_static("PLACES", &DsgLayers::PLACES)
       .def_readonly_static("MESH_PLACES", &DsgLayers::MESH_PLACES)
       .def_readonly_static("ROOMS", &DsgLayers::ROOMS)
-      .def_readonly_static("BUILDINGS", &DsgLayers::BUILDINGS);
+      .def_readonly_static("BUILDINGS", &DsgLayers::BUILDINGS)
+      .def_static(
+          "name_to_layer_id",
+          [](const std::string& name) -> std::optional<LayerId> {
+              return DsgLayers::nameToLayerId(name);
+          },
+          "name"_a);
 
   py::class_<LayerKey>(module, "LayerKey")
       .def(py::init<LayerId>())
@@ -745,6 +751,13 @@ PYBIND11_MODULE(_dsg_bindings, module) {
           "name"_a,
           "partition"_a,
           py::return_value_policy::reference_internal)
+      .def(
+          "remove_layer",
+          [](DynamicSceneGraph& graph, LayerId layer, PythonPartitionId partition) {
+              graph.removeLayer(layer, partition);
+          },
+          "layer"_a,
+          "partition"_a = 0)
       .def(
           "add_node",
           [](DynamicSceneGraph& graph,

diff --git a/python/src/spark_dsg/open3d_visualization.py b/python/src/spark_dsg/open3d_visualization.py
@@ -53,25 +53,31 @@
 
 
 LAYER_OFFSETS = {
-    DsgLayers.OBJECTS: 0.0,
-    DsgLayers.PLACES: 6.0,
-    DsgLayers.ROOMS: 12.0,
-    DsgLayers.BUILDINGS: 16.0,
+    DsgLayers.name_to_layer_id(DsgLayers.OBJECTS): 0.0,
+    DsgLayers.name_to_layer_id(DsgLayers.PLACES): 6.0,
+    DsgLayers.name_to_layer_id(DsgLayers.ROOMS): 12.0,
+    DsgLayers.name_to_layer_id(DsgLayers.BUILDINGS): 16.0,
 }
 
-LAYER_IDS = (DsgLayers.OBJECTS, DsgLayers.PLACES, DsgLayers.ROOMS, DsgLayers.BUILDINGS)
+LAYER_IDS = (
+    DsgLayers.name_to_layer_id(DsgLayers.OBJECTS),
+    DsgLayers.name_to_layer_id(DsgLayers.PLACES),
+    DsgLayers.name_to_layer_id(DsgLayers.ROOMS),
+    DsgLayers.name_to_layer_id(DsgLayers.BUILDINGS),
+)
+
 CATEGORY_MAP = {
-    "O": DsgLayers.OBJECTS,
-    "p": DsgLayers.PLACES,
-    "R": DsgLayers.ROOMS,
-    "B": DsgLayers.BUILDINGS,
+    "O": DsgLayers.name_to_layer_id(DsgLayers.OBJECTS),
+    "p": DsgLayers.name_to_layer_id(DsgLayers.PLACES),
+    "R": DsgLayers.name_to_layer_id(DsgLayers.ROOMS),
+    "B": DsgLayers.name_to_layer_id(DsgLayers.BUILDINGS),
 }
 
 LAYER_NAMES = {
-    DsgLayers.OBJECTS: "objects",
-    DsgLayers.PLACES: "places",
-    DsgLayers.ROOMS: "rooms",
-    DsgLayers.BUILDINGS: "buildings",
+    DsgLayers.name_to_layer_id(DsgLayers.OBJECTS): "objects",
+    DsgLayers.name_to_layer_id(DsgLayers.PLACES): "places",
+    DsgLayers.name_to_layer_id(DsgLayers.ROOMS): "rooms",
+    DsgLayers.name_to_layer_id(DsgLayers.BUILDINGS): "buildings",
 }
 
 
@@ -162,7 +168,7 @@ def _update_geometries(self, G):
 
         if self._include_dynamic:
             if G.has_layer(DsgLayers.AGENTS, "a"):
-                self._update_dynamic_layer(G.get_dynamic_layer(DsgLayers.AGENTS, "a"))
+                self._update_dynamic_layer(G.get_layer(DsgLayers.AGENTS, "a"))
 
     def _update_mesh_geometry(self, G):
         if not G.has_mesh():
@@ -334,7 +340,9 @@ def _update_names(self, G):
         self._names = [None] * len(self._id_map)
         for node_id in self._id_map:
             node = G.get_node(node_id)
-            if node.layer == DsgLayers.OBJECTS or node.layer == DsgLayers.ROOMS:
+            if node.layer == DsgLayers.name_to_layer_id(
+                DsgLayers.OBJECTS
+            ) or node.layer == DsgLayers.name_to_layer_id(DsgLayers.ROOMS):
                 self._names[self._id_map[node_id]] = node.attributes.name
 
         for name, point in zip(self._names, self._points):

diff --git a/src/dynamic_scene_graph.cpp b/src/dynamic_scene_graph.cpp
@@ -200,6 +200,12 @@ const Layer& DynamicSceneGraph::addLayer(const std::string& name,
 
 void DynamicSceneGraph::removeLayer(LayerId layer, PartitionId partition) {
   if (!partition) {
+    auto it = layers_.find(layer);
+    if (it != layers_.end()) {
+      while (!it->second->nodes_.empty()) {
+        removeNode(it->second->nodes_.begin()->first);
+      }
+    }
     layers_.erase(layer);
   }
 
@@ -208,7 +214,14 @@ void DynamicSceneGraph::removeLayer(LayerId layer, PartitionId partition) {
     return;
   }
 
-  iter->second.erase(partition);
+  auto part_iter = iter->second.find(partition);
+  if (part_iter != iter->second.end()) {
+    while (!part_iter->second->nodes_.empty()) {
+      removeNode(part_iter->second->nodes_.begin()->first);
+    }
+    iter->second.erase(part_iter);
+  }
+
   if (iter->second.empty()) {
     layer_partitions_.erase(iter);
   }