Fix/repart 2d

.github/workflows/ci.yaml

@@ -0,0 +1,65 @@
+---
+# based on https://github.com/actions/starter-workflows/blob/main/ci/cmake-single-platform.yml
+name: 'Spark-DSG Build and Test'
+on: [pull_request]
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Dependencies
+        run: sudo apt-get update && sudo apt install pipx
+      - name: Lint
+        run: pipx install pre-commit && cd ${{github.workspace}} && pre-commit run --all-files
+  cmake:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Dependencies
+        run: sudo apt-get update && sudo apt install libgtest-dev libeigen3-dev nlohmann-json3-dev libzmq3-dev
+      - name: Configure
+        run: cmake -B ${{github.workspace}}/build -DBUILD_TESTING=ON -DCMAKE_BUILD_TYPE=Release -DSPARK_DSG_BUILD_EXAMPLES=ON
+      - name: Build
+        run: cmake --build ${{github.workspace}}/build --config Release
+      - name: Test
+        working-directory: ${{github.workspace}}/build
+        run: ctest -C Release
+  python:
+    runs-on: ubuntu-latest
+    strategy: {matrix: {python-version: ['3.8', '3.10', '3.12']}}
+    steps:
+      - uses: actions/checkout@v4
+      - name: Setup Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Dependencies
+        run: |
+          python -m pip install --upgrade pip
+          sudo apt-get update && sudo apt install libgtest-dev libeigen3-dev nlohmann-json3-dev libzmq3-dev python3-dev
+      - name: Install
+        run: python -m pip install ${{github.workspace}}
+      - name: Test
+        run: |-
+          pip install torch --index-url https://download.pytorch.org/whl/cpu
+          pip install pytest networkx torch_geometric
+          pytest ${{github.workspace}}/python/tests
+  ros1:
+    runs-on: ubuntu-latest
+    container: ros:noetic-ros-base-focal
+    steps:
+      - uses: actions/checkout@v4
+        with: {path: src/spark_dsg}
+      - name: Dependencies
+        run: |
+          apt update && apt install -y python3-wstool python3-catkin-tools python3-vcstool git
+          rosdep update --rosdistro noetic && rosdep install --rosdistro noetic --from-paths src --ignore-src -r -y
+      - name: Configure
+        run: |
+          catkin init
+          catkin config --extend /opt/ros/noetic
+          catkin config -DCMAKE_BUILD_TYPE=Release -DSPARK_DSG_BUILD_EXAMPLES=ON
+      - name: Build
+        run: catkin build spark_dsg
+      - name: Test
+        run: catkin test spark_dsg

python/spark_dsg/mp3d.py

@@ -35,6 +35,7 @@
 """Functions for parsing a house file."""
 
 import heapq
+from typing import Any, Dict, List
 
 import numpy as np
 import seaborn as sns
@@ -192,7 +193,7 @@ def __init__(self, index, region, vertices, angle_deg=0.0):
             rotated_vertices = rotated_vertices[:-1]
             self._polygon_xy = shapely.geometry.Polygon(rotated_vertices)
 
-    def pos_on_same_floor(self, pos):
+    def pos_on_same_floor(self, pos, z_eps=0.0):
         """
         Check if a 3d position is within z-axis bounds.
 
@@ -202,23 +203,30 @@ def pos_on_same_floor(self, pos):
         Returns:
             bool: True if position falls inside [min_z, max_z]
         """
-        return self._min_z <= pos[2] <= self._max_z
+        return (self._min_z - z_eps) <= pos[2] <= (self._max_z + z_eps)
 
-    def pos_inside_room(self, pos):
+    def pos_inside_room(self, pos, eps=0.0, z_eps=0.0):
         """
-        Check if a 3d position falls within the bounds of the room.
+        Check if a 3d position falls within the bound of the room
+        (i.e. within eps bound in x-y plane and z_eps bound along z-axis).
 
         Args:
             pos (List[float]): 3d position to check
 
-        Returns:
-            bool: True if position falls inside [min_z, max_z] and polygon bounds
+        Returns:eps
+            bool: True if position falls inside room 2D polygon and vertical height
+            float: x-y distance to the room polygon or "inf" if not on the same floor
         """
-        if not self.pos_on_same_floor(pos):
-            return False
+        if not self.pos_on_same_floor(pos, z_eps=z_eps):
+            return False, float("inf")
 
         xy_pos = shapely.geometry.Point(pos[0], pos[1])
-        return self._polygon_xy.contains(xy_pos)
+
+        if self._polygon_xy.contains(xy_pos):
+            return True, 0.0
+
+        dist = self._polygon_xy.distance(xy_pos)
+        return dist < eps, dist
 
     def get_polygon_xy(self):
         """Get the xy bounding polygon of the room."""
@@ -254,15 +262,15 @@ def semantic_label(self):
         return ord(self._label)
 
 
-def load_mp3d_info(house_path):
+def load_mp3d_info(house_path) -> Dict[str, List[Dict[str, Any]]]:
     """
     Load room info from a GT house file.
 
     Args:
         house_path (str): Path to house file
 
     Returns:
-        Dict[str, List[Dict[Str, Any]]]): Parsed house file information
+        Dict[str, List[Dict[Str, Any]]]: Parsed house file information
     """
     info = {x: [] for x in PARSERS}
     with open(house_path, "r") as fin:
@@ -277,7 +285,7 @@ def load_mp3d_info(house_path):
     return info
 
 
-def get_rooms_from_mp3d_info(mp3d_info, angle_deg=-90.0):
+def get_rooms_from_mp3d_info(mp3d_info, angle_deg=-90.0) -> List[Mp3dRoom]:
     """
     Generate a list of Mp3dRoom objects from ground-truth segmentation.
 
@@ -393,6 +401,8 @@ def repartition_rooms(
     min_iou_threshold=0.0,
     colors=None,
     verbose=False,
+    eps=0.0,
+    z_eps=0.0,
 ):
     """
     Create a copy of the DSG with ground-truth room nodes.
@@ -439,16 +449,25 @@ def repartition_rooms(
     missing_nodes = []
     for place in G.get_layer(DsgLayers.PLACES).nodes:
         pos = G.get_position(place.id.value)
+        best_room = None
+        best_error = float("inf")
 
         for room in new_rooms:
-            if not room.pos_inside_room(pos):
-                continue
-
-            room_id = room.get_id()
+            inside, error = room.pos_inside_room(pos, eps=eps, z_eps=z_eps)
+            if error < best_error:
+                best_error = error
+                if inside:  # update best_room if inside
+                    best_room = room
+
+        if best_room is not None:
+            room_id = best_room.get_id()
             G.insert_edge(place.id.value, room_id.value)
-            break  # avoid labeling node as missing
-        else:
+        elif verbose:
             missing_nodes.append(place)
+            print(
+                f"Place node {place.id.value} not assigned to any room "
+                f"(min dist to rooms on the same floor: {best_error:.2f})."
+            )
     if verbose:
         print(f"Found {len(missing_nodes)} places node outside of room segmentations.")
 

python/spark_dsg/visualization.py

@@ -39,16 +39,23 @@
 import numpy as np
 
 # DSG plot style
-NODE_TYPE_OFFSET = {"B": 30, "R": 25, "p": 10, "O": 0}
+LAYER_TYPE_OFFSET = {
+    5: 30,  # buildings
+    4: 25,  # rooms
+    3: 10,  # places
+    2: 0,  # objects
+}
 
-NODE_TYPE_TO_COLOR = {"B": "#636EFA", "R": "#EF553B", "p": "#AB63FA", "O": "#00CC96"}
+LAYER_TYPE_TO_COLOR = {
+    5: "#636EFA",  # buildings
+    4: "#EF553B",  # rooms
+    3: "#AB63FA",  # places
+    2: "#00CC96",  # objects
+}
 
 
-def z_offset(node) -> np.ndarray:
-    """Take a node and returns an offset in the z direction according to node type."""
-    offset = node.attributes.position.copy()
-    offset[2] += NODE_TYPE_OFFSET[node.id.category]
-    return offset
+def z_offset(layer_idx: int) -> np.ndarray:
+    return np.array([0, 0, LAYER_TYPE_OFFSET[layer_idx]])
 
 
 def _draw_layer_nodes(
@@ -65,17 +72,19 @@ def _draw_layer_nodes(
     pos, colors, text = [], [], []
 
     for node in layer.nodes:
-        pos.append(np.squeeze(z_offset(node)))
+        pos.append(np.array(node.attributes.position) + z_offset(layer.id))
         if color_func is None:
-            colors.append(NODE_TYPE_TO_COLOR[node.id.category])
+            colors.append(LAYER_TYPE_TO_COLOR[layer.id])
         else:
             colors.append(color_func(node))
-
         if text_func is None:
             text.append(str(node.id))
         else:
             text.append(text_func(node))
 
+    if len(pos) == 0:
+        return
+
     pos = np.array(pos)
     fig.add_trace(
         go.Scatter3d(
@@ -105,8 +114,13 @@ def plot_scene_graph(G, title=None, figure_path=None, layer_settings=None):
     for layer in G.layers:
         has_settings = layer_settings is not None and layer.id in layer_settings
         settings = {} if not has_settings else layer_settings[layer.id]
-        print(settings)
-        _draw_layer_nodes(fig, layer, **settings)
+        if layer.id in LAYER_TYPE_OFFSET.keys():
+            _draw_layer_nodes(fig, layer, **settings)
+        else:
+            print(
+                f"Skipping layer (id={layer.id}), "
+                f"containing {len([n for n in layer.nodes])} nodes."
+            )
 
     # edges
     x_lines, y_lines, z_lines = [], [], []
@@ -116,20 +130,25 @@ def plot_scene_graph(G, title=None, figure_path=None, layer_settings=None):
         source = G.get_node(edge.source)
         target = G.get_node(edge.target)
 
-        start_offset = z_offset(source)
-        end_offset = z_offset(target)
+        if (
+            source.layer.layer not in LAYER_TYPE_OFFSET.keys()
+            or target.layer.layer not in LAYER_TYPE_OFFSET.keys()
+        ):
+            continue
+        source_pos = np.array(source.attributes.position) + z_offset(source.layer.layer)
+        target_pos = np.array(target.attributes.position) + z_offset(target.layer.layer)
 
         # intralayer edges
         if source.layer == target.layer:
-            x_lines_dark += [start_offset[0], end_offset[0], None]
-            y_lines_dark += [start_offset[1], end_offset[1], None]
-            z_lines_dark += [start_offset[2], end_offset[2], None]
+            x_lines_dark += [source_pos[0], target_pos[0], None]
+            y_lines_dark += [source_pos[1], target_pos[1], None]
+            z_lines_dark += [source_pos[2], target_pos[2], None]
 
         # interlayer edges
         else:
-            x_lines += [start_offset[0], end_offset[0], None]
-            y_lines += [start_offset[1], end_offset[1], None]
-            z_lines += [start_offset[2], end_offset[2], None]
+            x_lines += [source_pos[0], target_pos[0], None]
+            y_lines += [source_pos[1], target_pos[1], None]
+            z_lines += [source_pos[2], target_pos[2], None]
 
     # add interlayer edges to plot
     fig.add_trace(