add hash for layer key and skip partitioned nodes in torch conversion

python/src/spark_dsg/__init__.py

@@ -39,7 +39,7 @@
 from spark_dsg._dsg_bindings import *
 from spark_dsg._dsg_bindings import (BoundingBoxType, DsgLayers,
                                      DynamicSceneGraph, EdgeAttributes,
-                                     LayerView, NodeAttributes,
+                                     LayerKey, LayerView, NodeAttributes,
                                      SceneGraphLayer,
                                      compute_ancestor_bounding_box)
 from spark_dsg.open3d_visualization import render_to_open3d
@@ -110,6 +110,12 @@ def _add_metadata_interface(obj):
     obj.add_metadata = _add_metadata
 
 
+def _hash_layerkey(key):
+    return hash((key.layer, key.partition))
+
+
+LayerKey.__hash__ = _hash_layerkey
+
 _add_metadata_interface(DynamicSceneGraph)
 _add_metadata_interface(NodeAttributes)
 _add_metadata_interface(EdgeAttributes)

python/src/spark_dsg/torch_conversion.py

@@ -325,6 +325,9 @@ def scene_graph_to_torch_heterogeneous(
     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] = []
@@ -356,6 +359,9 @@ def scene_graph_to_torch_heterogeneous(
     for edge in G.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]
         if edge_type not in edge_indices:
             edge_indices[edge_type] = []