basically functional manual Room annotation

dcist_launch_system/scripts/gt_room_adder.py

@@ -0,0 +1,96 @@
+import os
+from typing import Optional
+
+import numpy as np
+import spark_dsg
+import yaml
+from scipy.spatial.transform import Rotation as Rot
+
+
+class RoomExtents:
+    def __init__(self, path_to_yaml: Optional[str] = None):
+        self.boxes: list[list[spark_dsg.BoundingBox]] = []
+
+        if path_to_yaml is None or not os.path.exists(path_to_yaml):
+            raise Exception(
+                "Currently RoomExtents must be initialized from a valid yaml path"
+            )
+
+        with open(path_to_yaml, "r") as f:
+            root = yaml.safe_load(f)
+
+        for key in sorted(root.keys(), key=lambda k: int(k)):
+            group = []
+            for box_node in root[key]:
+                center = np.array(box_node["center"], dtype=float)
+                dimensions = np.array(box_node["extents"], dtype=float)
+                rot_dict = box_node["rotation"]
+                rotation = np.array(
+                    [rot_dict["x"], rot_dict["y"], rot_dict["z"], rot_dict["w"]],
+                    dtype=float,
+                )
+
+                rotmat = Rot.from_quat(rotation).as_matrix()
+                group.append(
+                    spark_dsg.BoundingBox(
+                        spark_dsg.BoundingBoxType.AABB, dimensions, center, rotmat
+                    )
+                )
+
+            self.boxes.append(group)
+
+
+def add_gt_rooms(
+    dsg: spark_dsg.DynamicSceneGraph,
+    room_boxes: list[list[spark_dsg.BoundingBox]],
+    semantic_labels: list[int],
+):
+    for idx, boxes in enumerate(room_boxes):
+        box_centers = np.mean([b.world_P_center for b in boxes], axis=0)
+        room_position = box_centers
+        room_attrs = spark_dsg.RoomNodeAttributes()
+        room_attrs.semantic_label = semantic_labels[idx]
+        room_attrs.position = room_position
+        dsg.add_node(
+            spark_dsg.DsgLayers.ROOMS, spark_dsg.NodeSymbol("R", idx), room_attrs
+        )
+
+    for idx, boxes in enumerate(room_boxes):
+        print(f"{idx}, {boxes}")
+        for box in boxes:
+            for node in dsg.get_layer(spark_dsg.DsgLayers.MESH_PLACES).nodes:
+                pos = node.attributes.position
+                if box.contains(pos):
+                    print(f"R{idx} -> {node.id}")
+                    dsg.insert_edge(spark_dsg.NodeSymbol("R", idx), node.id)
+
+
+def add_manual_connections(dsg, connections):
+    for edge in connections:
+        s = edge[0]
+        t = edge[1]
+        dsg.insert_edge(spark_dsg.NodeSymbol("R", s), spark_dsg.NodeSymbol("R", t))
+
+
+if __name__ == "__main__":
+    dsg_path = "/home/aaron/adt4_output/wednesday_afternoon_1_processed/hydra/backend/dsg_with_mesh.json"
+    G = spark_dsg.DynamicSceneGraph.load(dsg_path)
+
+    manual_room_connections = [(0, 1), (1, 2), (2, 3)]
+    bb_path = "gt_room_bounding_boxes.yaml"
+
+    extents = RoomExtents(bb_path)
+
+    semantic_labels = [2, 0, 1]
+
+    add_gt_rooms(G, extents.boxes, semantic_labels)
+    add_manual_connections(G, manual_room_connections)
+
+    for node in G.get_layer(spark_dsg.DsgLayers.ROOMS).nodes:
+        print(node.id)
+        print(node.attributes)
+
+    for edge in G.get_layer(spark_dsg.DsgLayers.ROOMS).edges:
+        print(edge)
+
+    G.save("test_dsg.json")

dcist_launch_system/scripts/gt_room_creator.py

@@ -0,0 +1,125 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import trimesh
+import yaml
+
+DEFAULT_Z = 2.5
+DEFAULT_HEIGHT = 8.0
+DEFAULT_ROTATION = {"w": 1.0, "x": 0.0, "y": 0.0, "z": 0.0}
+
+
+bounding_boxes = {}  # group_id -> list of boxes
+clicks = []
+current_group = 1
+
+
+def onclick(event):
+    if event.xdata is None or event.ydata is None:
+        return
+    clicks.append((event.xdata, event.ydata))
+    print(f"[Group {current_group}] Click: {event.xdata:.2f}, {event.ydata:.2f}")
+
+    if len(clicks) == 2:
+        x1, y1 = clicks[0]
+        x2, y2 = clicks[1]
+        center = [(x1 + x2) / 2, (y1 + y2) / 2, DEFAULT_Z]
+        extents = [abs(x2 - x1), abs(y2 - y1), DEFAULT_HEIGHT]
+        print("centers: ", center)
+        print(type(round(center[0], 2)))
+        box = {
+            "center": [float(round(x, 2)) for x in center],
+            "extents": [float(round(x, 2)) for x in extents],
+            "rotation": DEFAULT_ROTATION.copy(),
+        }
+
+        if current_group not in bounding_boxes:
+            bounding_boxes[current_group] = []
+        bounding_boxes[current_group].append(box)
+        print(f"  → Added to group {current_group}: {box}")
+        print("bounding_boxes: ", bounding_boxes)
+
+        ax.add_patch(
+            plt.Rectangle(
+                (min(x1, x2), min(y1, y2)),
+                abs(x2 - x1),
+                abs(y2 - y1),
+                fill=True,
+                edgecolor="red",
+                linewidth=2,
+            )
+        )
+        ax.text(
+            center[0],
+            center[1],
+            str(current_group),
+            color="blue",
+            ha="center",
+            va="center",
+        )
+        plt.draw()
+        clicks.clear()
+
+
+def onkey(event):
+    global current_group
+
+    if event.key == "enter":
+        with open("gt_room_bounding_boxes.yaml", "w") as f:
+            yaml.dump(bounding_boxes, f, sort_keys=True)
+        print("✅ Saved bounding_boxes.yaml")
+        plt.close()
+    elif event.key.isdigit() and int(event.key) > 0:
+        current_group = int(event.key)
+        print(f"🔀 Switched to group {current_group}")
+
+
+# Load the .ply file
+mesh_fn = (
+    "/home/aaron/adt4_output/wednesday_afternoon_1_processed/hydra/backend/mesh.ply"
+)
+mesh = trimesh.load(mesh_fn)
+
+# Extract vertex positions (Nx3)
+vertices = mesh.vertices
+
+# Choose a projection: XY, XZ, or YZ
+proj = "xy"  # change to "xz" or "yz" if desired
+
+if proj == "xy":
+    x = vertices[:, 0]
+    y = vertices[:, 1]
+elif proj == "xz":
+    x = vertices[:, 0]
+    y = vertices[:, 2]
+elif proj == "yz":
+    x = vertices[:, 1]
+    y = vertices[:, 2]
+else:
+    raise ValueError("Invalid projection: choose from 'xy', 'xz', 'yz'")
+
+# Plot
+fig, ax = plt.subplots()
+plt.xlabel(proj[0].upper())
+plt.ylabel(proj[1].upper())
+plt.title(f"{proj.upper()} Projection of PLY Vertices")
+
+
+ax.set_title("Click 2 corners per box. Press 1–9 to change group. Enter to finish.")
+fig.canvas.mpl_connect("button_press_event", onclick)
+fig.canvas.mpl_connect("key_press_event", onkey)
+
+
+# plt.axis("equal")
+# plt.grid(True)
+
+rs = mesh.visual.vertex_colors[:, 0]
+gs = mesh.visual.vertex_colors[:, 1]
+bs = mesh.visual.vertex_colors[:, 2]
+colors = np.array([rs, gs, bs])
+colors = mesh.visual.vertex_colors[:, :4]
+colors[:, 3] = 100
+print("color shape: ", colors.shape)
+
+plt.scatter(x, y, s=2, c=colors / 255)
+
+plt.show()