add unpartitioned iterators and fix torch graph conversions

include/spark_dsg/dynamic_scene_graph.h

@@ -446,6 +446,13 @@ class DynamicSceneGraph {
    */
   std::vector<EdgeKey> getNewEdges(bool clear_new = false);
 
+  /**
+   * @brief Get whether an edge points to a partition
+   * @param edge Edge to check
+   * @return True if at least one end point is in a partition
+   */
+  bool edgeToPartition(const SceneGraphEdge& edge) const;
+
   /**
    * @brief Track which edges get used during a serialization update
    */

python/bindings/include/spark_dsg/python/python_layer_view.h

@@ -93,20 +93,21 @@ class PartitionIter {
 
 class GlobalLayerIter {
  public:
-  GlobalLayerIter(const DynamicSceneGraph& graph);
+  GlobalLayerIter(const DynamicSceneGraph& graph, bool include_partitions = true);
   LayerView operator*() const;
   GlobalLayerIter& operator++();
   bool operator==(const IterSentinel&) const;
   bool operator!=(const IterSentinel&) const { return !(*this == IterSentinel()); }
 
  private:
+  bool include_partitions_;
   LayerIter layers_;
   PartitionIter partitions_;
 };
 
 class GlobalNodeIter {
  public:
-  GlobalNodeIter(const DynamicSceneGraph& dsg);
+  GlobalNodeIter(const DynamicSceneGraph& dsg, bool include_partitions = true);
   void setNodeIter();
   const SceneGraphNode* operator*() const;
   GlobalNodeIter& operator++();
@@ -120,14 +121,17 @@ class GlobalNodeIter {
 
 class GlobalEdgeIter {
  public:
-  GlobalEdgeIter(const DynamicSceneGraph& dsg);
+  GlobalEdgeIter(const DynamicSceneGraph& dsg, bool include_partitions = true);
   const SceneGraphEdge* operator*() const;
   void setEdgeIter();
+  void findNextValidEdge();
   GlobalEdgeIter& operator++();
   bool operator==(const IterSentinel&);
 
  private:
+  bool include_partitions_;
   bool started_interlayer_;
+  const DynamicSceneGraph& dsg_;
   GlobalLayerIter layers_;
   EdgeIter curr_edge_iter_;
   EdgeIter interlayer_edge_iter_;

python/bindings/include/spark_dsg/python/scene_graph_iterators.h

@@ -45,7 +45,8 @@ class NodeIter {
   NodeIter(const SceneGraphLayer::Nodes& container);
   const SceneGraphNode* operator*() const;
   NodeIter& operator++();
-  bool operator==(const IterSentinel&);
+  bool operator==(const IterSentinel&) const;
+  bool operator!=(const IterSentinel&) const { return !(*this == IterSentinel()); }
 
  private:
   bool valid_;
@@ -55,11 +56,12 @@ class NodeIter {
 
 class EdgeIter {
  public:
-   EdgeIter();
+  EdgeIter();
   EdgeIter(const SceneGraphLayer::Edges& container);
   const SceneGraphEdge* operator*() const;
   EdgeIter& operator++();
-  bool operator==(const IterSentinel&);
+  bool operator==(const IterSentinel&) const;
+  bool operator!=(const IterSentinel&) const { return !(*this == IterSentinel()); }
 
  private:
   bool valid_;

python/bindings/src/python_layer_view.cpp

@@ -133,15 +133,18 @@ bool PartitionIter::operator==(const IterSentinel&) const {
   return curr_layer_iter_ == end_layer_iter_ && curr_iter_ == end_iter_;
 }
 
-GlobalLayerIter::GlobalLayerIter(const DynamicSceneGraph& graph)
-    : layers_(graph.layers()), partitions_(graph.layer_partitions()) {}
+GlobalLayerIter::GlobalLayerIter(const DynamicSceneGraph& graph,
+                                 bool include_partitions)
+    : include_partitions_(include_partitions),
+      layers_(graph.layers()),
+      partitions_(graph.layer_partitions()) {}
 
 LayerView GlobalLayerIter::operator*() const {
   if (layers_ != IterSentinel()) {
     return *layers_;
   }
 
-  if (partitions_ != IterSentinel()) {
+  if (include_partitions_ && partitions_ != IterSentinel()) {
     return *partitions_;
   }
 
@@ -151,19 +154,20 @@ LayerView GlobalLayerIter::operator*() const {
 GlobalLayerIter& GlobalLayerIter::operator++() {
   if (layers_ != IterSentinel()) {
     ++layers_;
-  } else if (partitions_ != IterSentinel()) {
+  } else if (include_partitions_ && partitions_ != IterSentinel()) {
     ++partitions_;
   }
 
   return *this;
 }
 
 bool GlobalLayerIter::operator==(const IterSentinel&) const {
-  return layers_ == IterSentinel() && partitions_ == IterSentinel();
+  return layers_ == IterSentinel() &&
+         (!include_partitions_ || partitions_ == IterSentinel());
 }
 
-GlobalNodeIter::GlobalNodeIter(const DynamicSceneGraph& dsg)
-    : valid_(true), layers_(dsg) {
+GlobalNodeIter::GlobalNodeIter(const DynamicSceneGraph& dsg, bool include_partitions)
+    : valid_(true), layers_(dsg, include_partitions) {
   setNodeIter();
 }
 
@@ -205,9 +209,11 @@ bool GlobalNodeIter::operator==(const IterSentinel&) {
   return curr_node_iter_ == IterSentinel() && layers_ == IterSentinel();
 }
 
-GlobalEdgeIter::GlobalEdgeIter(const DynamicSceneGraph& dsg)
-    : started_interlayer_(false),
-      layers_(dsg),
+GlobalEdgeIter::GlobalEdgeIter(const DynamicSceneGraph& dsg, bool include_partitions)
+    : include_partitions_(include_partitions),
+      started_interlayer_(false),
+      dsg_(dsg),
+      layers_(dsg, include_partitions),
       interlayer_edge_iter_(dsg.interlayer_edges()) {
   setEdgeIter();
 }
@@ -216,13 +222,27 @@ const SceneGraphEdge* GlobalEdgeIter::operator*() const {
   return started_interlayer_ ? *interlayer_edge_iter_ : *curr_edge_iter_;
 }
 
+void GlobalEdgeIter::findNextValidEdge() {
+  if (include_partitions_) {
+    // every edge is valid if we include partitions
+    return;
+  }
+
+  while (dsg_.edgeToPartition(*(*interlayer_edge_iter_)) &&
+         interlayer_edge_iter_ != IterSentinel()) {
+    ++interlayer_edge_iter_;
+  }
+}
+
 void GlobalEdgeIter::setEdgeIter() {
   if (started_interlayer_ || layers_ == IterSentinel()) {
     started_interlayer_ = true;
+    findNextValidEdge();
     return;
   }
 
   curr_edge_iter_ = (*layers_).edges();
+
   while (curr_edge_iter_ == IterSentinel()) {
     ++layers_;
     if (layers_ == IterSentinel()) {
@@ -241,6 +261,7 @@ GlobalEdgeIter& GlobalEdgeIter::operator++() {
 
   if (started_interlayer_) {
     ++interlayer_edge_iter_;
+    findNextValidEdge();
     return *this;
   }
 

python/bindings/src/scene_graph_iterators.cpp

@@ -58,7 +58,7 @@ NodeIter& NodeIter::operator++() {
   return *this;
 }
 
-bool NodeIter::operator==(const IterSentinel&) {
+bool NodeIter::operator==(const IterSentinel&) const {
   return !valid_ || curr_iter_ == end_iter_;
 }
 
@@ -84,7 +84,7 @@ EdgeIter& EdgeIter::operator++() {
   return *this;
 }
 
-bool EdgeIter::operator==(const IterSentinel&) {
+bool EdgeIter::operator==(const IterSentinel&) const {
   return !valid_ || curr_iter_ == end_iter_;
 }
 

python/bindings/src/spark_dsg_bindings.cpp

@@ -894,6 +894,20 @@ PYBIND11_MODULE(_dsg_bindings, module) {
           },
           nullptr,
           py::return_value_policy::reference_internal)
+      .def_property(
+          "unpartitioned_nodes",
+          [](const DynamicSceneGraph& graph) {
+            return py::make_iterator(GlobalNodeIter(graph, false), IterSentinel());
+          },
+          nullptr,
+          py::return_value_policy::reference_internal)
+      .def_property(
+          "unpartitioned_edges",
+          [](const DynamicSceneGraph& graph) {
+            return py::make_iterator(GlobalEdgeIter(graph, false), IterSentinel());
+          },
+          nullptr,
+          py::return_value_policy::reference_internal)
       .def_property(
           "interlayer_edges",
           [](const DynamicSceneGraph& graph) {

python/src/spark_dsg/torch_conversion.py

@@ -44,20 +44,15 @@
 
 import numpy as np
 
-from spark_dsg._dsg_bindings import (DsgLayers, DynamicSceneGraph, LayerView,
+from spark_dsg._dsg_bindings import (DynamicSceneGraph, LayerView,
                                      SceneGraphEdge, SceneGraphLayer,
                                      SceneGraphNode)
 
 NodeConversionFunc = Callable[[DynamicSceneGraph, SceneGraphNode], np.ndarray]
 EdgeConversionFunc = Callable[[DynamicSceneGraph, SceneGraphEdge], np.ndarray]
 
 
-DEFAULT_LAYER_MAP = {
-    DsgLayers.OBJECTS: "objects",
-    DsgLayers.PLACES: "places",
-    DsgLayers.ROOMS: "rooms",
-    DsgLayers.BUILDINGS: "buildings",
-}
+DEFAULT_LAYER_MAP = {2: "objects", 3: "places", 4: "rooms", 5: "buildings"}
 
 
 def _centroid_bbx_embedding(G, x) -> NodeConversionFunc:
@@ -224,9 +219,9 @@ def scene_graph_to_torch_homogeneous(
     node_ids = []
     id_map = {}
 
-    for node in G.nodes:
+    for node in G.unpartitioned_nodes:
         idx = len(node_features)
-        node_masks[node.layer][idx] = True
+        node_masks[node.layer.layer][idx] = True
         node_positions[idx, :] = torch.tensor(
             np.squeeze(node.attributes.position), dtype=dtype_float
         )
@@ -241,7 +236,7 @@ def scene_graph_to_torch_homogeneous(
 
     edge_index = torch.zeros((2, M), dtype=torch.int64)
     edge_features = []
-    for idx, edge in enumerate(G.edges):
+    for idx, edge in enumerate(G.unpartitioned_edges):
         edge_index[:, idx] = torch.tensor(_get_directed_edge(G, edge, id_map))
 
         if edge_converter is not None:
@@ -324,21 +319,19 @@ def scene_graph_to_torch_heterogeneous(
     node_ids = {}
     id_map = {}
 
-    for node in G.nodes:
-        if node.layer.partition != 0:
-            continue
-
-        if node.layer not in node_features:
-            node_features[node.layer] = []
-            node_positions[node.layer] = []
-            node_labels[node.layer] = []
-            node_ids[node.layer] = []
-
-        idx = len(node_features[node.layer])
-        node_positions[node.layer].append(np.squeeze(node.attributes.position))
-        node_features[node.layer].append(node_converter(G, node))
-        node_labels[node.layer].append(node.attributes.semantic_label)
-        node_ids[node.layer].append(node.id.value)
+    for node in G.unpartitioned_nodes:
+        layer_id = node.layer.layer
+        if layer_id not in node_features:
+            node_features[layer_id] = []
+            node_positions[layer_id] = []
+            node_labels[layer_id] = []
+            node_ids[layer_id] = []
+
+        idx = len(node_features[layer_id])
+        node_positions[layer_id].append(np.squeeze(node.attributes.position))
+        node_features[layer_id].append(node_converter(G, node))
+        node_labels[layer_id].append(node.attributes.semantic_label)
+        node_ids[layer_id].append(node.id.value)
         id_map[node.id.value] = idx
 
     for layer in node_features:
@@ -356,13 +349,10 @@ def scene_graph_to_torch_heterogeneous(
 
     edge_indices = {}
     edge_features = {}
-    for edge in G.edges:
+    for edge in G.unpartitioned_edges:
         source = G.get_node(edge.source)
         target = G.get_node(edge.target)
-        if source.layer.partition != 0 or target.layer.partition != 0:
-            continue
-
-        edge_type = edge_map[source.layer][target.layer][1]
+        edge_type = edge_map[source.layer.layer][target.layer.layer][1]
         if edge_type not in edge_indices:
             edge_indices[edge_type] = []
             edge_features[edge_type] = []

src/dynamic_scene_graph.cpp

@@ -445,7 +445,6 @@ size_t DynamicSceneGraph::numUnpartitionedEdges() const {
   return total_edges;
 }
 
-
 bool DynamicSceneGraph::empty() const { return numNodes() == 0; }
 
 Eigen::Vector3d DynamicSceneGraph::getPosition(NodeId node_id) const {
@@ -576,6 +575,11 @@ std::vector<EdgeKey> DynamicSceneGraph::getNewEdges(bool clear_new) {
   return to_return;
 }
 
+bool DynamicSceneGraph::edgeToPartition(const SceneGraphEdge& edge) const {
+  const auto lookup = lookupEdge(edge.source, edge.target);
+  return lookup.source.partition || lookup.target.partition;
+}
+
 void DynamicSceneGraph::markEdgesAsStale() {
   for (auto& [layer_id, layer] : layers_) {
     layer->edges_.setStale();