From b48cbdee7b2e93e92e0dcee94c8a581ec022c4d3 Mon Sep 17 00:00:00 2001 From: Nathan Hughes Date: Fri, 22 Nov 2024 18:33:04 +0000 Subject: [PATCH] add unpartitioned iterators and fix torch graph conversions --- include/spark_dsg/dynamic_scene_graph.h | 7 +++ .../spark_dsg/python/python_layer_view.h | 10 ++-- .../spark_dsg/python/scene_graph_iterators.h | 8 +-- python/bindings/src/python_layer_view.cpp | 41 +++++++++++---- python/bindings/src/scene_graph_iterators.cpp | 4 +- python/bindings/src/spark_dsg_bindings.cpp | 14 ++++++ python/src/spark_dsg/torch_conversion.py | 50 ++++++++----------- src/dynamic_scene_graph.cpp | 6 ++- 8 files changed, 91 insertions(+), 49 deletions(-) diff --git a/include/spark_dsg/dynamic_scene_graph.h b/include/spark_dsg/dynamic_scene_graph.h index b8a4fa9..107d9c4 100644 --- a/include/spark_dsg/dynamic_scene_graph.h +++ b/include/spark_dsg/dynamic_scene_graph.h @@ -446,6 +446,13 @@ class DynamicSceneGraph { */ std::vector 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 */ diff --git a/python/bindings/include/spark_dsg/python/python_layer_view.h b/python/bindings/include/spark_dsg/python/python_layer_view.h index 9198be1..2651995 100644 --- a/python/bindings/include/spark_dsg/python/python_layer_view.h +++ b/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_; diff --git a/python/bindings/include/spark_dsg/python/scene_graph_iterators.h b/python/bindings/include/spark_dsg/python/scene_graph_iterators.h index f2a5910..ab67dcd 100644 --- a/python/bindings/include/spark_dsg/python/scene_graph_iterators.h +++ b/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_; diff --git a/python/bindings/src/python_layer_view.cpp b/python/bindings/src/python_layer_view.cpp index a205690..3c316c0 100644 --- a/python/bindings/src/python_layer_view.cpp +++ b/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,7 +154,7 @@ LayerView GlobalLayerIter::operator*() const { GlobalLayerIter& GlobalLayerIter::operator++() { if (layers_ != IterSentinel()) { ++layers_; - } else if (partitions_ != IterSentinel()) { + } else if (include_partitions_ && partitions_ != IterSentinel()) { ++partitions_; } @@ -159,11 +162,12 @@ GlobalLayerIter& GlobalLayerIter::operator++() { } 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; } diff --git a/python/bindings/src/scene_graph_iterators.cpp b/python/bindings/src/scene_graph_iterators.cpp index 487649d..3e9f46e 100644 --- a/python/bindings/src/scene_graph_iterators.cpp +++ b/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_; } diff --git a/python/bindings/src/spark_dsg_bindings.cpp b/python/bindings/src/spark_dsg_bindings.cpp index 4e25977..323a84a 100644 --- a/python/bindings/src/spark_dsg_bindings.cpp +++ b/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) { diff --git a/python/src/spark_dsg/torch_conversion.py b/python/src/spark_dsg/torch_conversion.py index 008ad1b..c652277 100644 --- a/python/src/spark_dsg/torch_conversion.py +++ b/python/src/spark_dsg/torch_conversion.py @@ -44,7 +44,7 @@ import numpy as np -from spark_dsg._dsg_bindings import (DsgLayers, DynamicSceneGraph, LayerView, +from spark_dsg._dsg_bindings import (DynamicSceneGraph, LayerView, SceneGraphEdge, SceneGraphLayer, SceneGraphNode) @@ -52,12 +52,7 @@ 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] = [] diff --git a/src/dynamic_scene_graph.cpp b/src/dynamic_scene_graph.cpp index 91890f5..9fac9d9 100644 --- a/src/dynamic_scene_graph.cpp +++ b/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 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();