import json import cv2 import networkx as nx import numpy as np from pathlib import Path import os import heapq import math from collections import deque import matplotlib.pyplot as plt from PIL import Image import random class BuildingGraph: def __init__(self, default_floor="UNKNOWN"): """ Initialize the BuildingGraph with an empty graph. Args: default_floor (str): Default floor id to use for nodes when not provided. Outdoors will always be stored as floor="NA". """ self.graph = nx.Graph() self.node_types = {"room": [], "door": [], "corridor": [], "outside": [], "transition": []} self.default_floor = default_floor # used if caller doesn't pass floor_id self.property = {} # --------------------------- Utilities: floors ----------------------------- def set_default_floor(self, floor_id: str): """Update the default floor used for new nodes (non-outdoor).""" self.default_floor = floor_id def _resolve_floor(self, node_type: str, floor_id): """ Decide which floor to store for this node. Outdoors are explicitly tagged as 'NA'. """ if node_type == "outside": return "NA" return self.default_floor if floor_id is None else floor_id # -------------------------------------------------------------------------- def add_node(self, node_id, node_type, position, pixels=None, floor_id=None): """ Add a node to the graph. Args: node_id (str): Unique identifier for the node. node_type (str): Type of the node ('room', 'door', 'corridor', 'outside'). position (tuple): (x, y) coordinates of the node. pixels (list, optional): List of pixels belonging to the node. floor_id (str, optional): Floor id for this node (overrides default). Outdoors will always be saved as 'NA'. """ if node_type not in self.node_types: raise ValueError( f"Invalid node type: {node_type}. Must be one of {list(self.node_types.keys())}." ) if pixels is None: pixels = [] node_floor = self._resolve_floor(node_type, floor_id) # Add node to the graph self.graph.add_node( node_id, type=node_type, position=position, pixels=pixels, floor=node_floor, ) # Update the node_types dictionary self.node_types[node_type].append(node_id) def _ensure_edge_metrics(self, overwrite_weight_if_one=True): """Backfill `distance` and (optionally) replace default weight 1 with distance.""" import math for u, v, ed in self.graph.edges(data=True): pos1 = self.graph.nodes[u].get("position") pos2 = self.graph.nodes[v].get("position") # compute distance if possible dist = ed.get("distance") if (dist is None) and (pos1 is not None) and (pos2 is not None): try: x1, y1 = float(pos1[0]), float(pos1[1]) x2, y2 = float(pos2[0]), float(pos2[1]) dist = ((x1 - x2)**2 + (y1 - y2)**2) ** 0.5 ed["distance"] = float(dist) except Exception: pass # upgrade weight if it's missing or the legacy default (1) if overwrite_weight_if_one: if ("weight" not in ed) or (ed["weight"] in (None, 1, 1.0)): if dist is not None: ed["weight"] = float(dist) def add_edge(self, node_id_1, node_id_2, weight=None): """ Add an edge and record BOTH `weight` and geometric `distance`. If `weight` is None, default to Euclidean distance (or 1.0 if positions missing). """ import math if not self.graph.has_node(node_id_1) or not self.graph.has_node(node_id_2): print(f"Warning: attempted to add edge between non-existent nodes " f"'{node_id_1}' and '{node_id_2}'. Skipping.") return pos1 = self.graph.nodes[node_id_1].get("position") pos2 = self.graph.nodes[node_id_2].get("position") distance = None if pos1 is not None and pos2 is not None: try: x1, y1 = float(pos1[0]), float(pos1[1]) x2, y2 = float(pos2[0]), float(pos2[1]) distance = ( (x1 - x2)**2 + (y1 - y2)**2 ) ** 0.5 except Exception: distance = None if weight is None: weight = float(distance) if distance is not None else 1.0 # write both attrs if self.graph.has_edge(node_id_1, node_id_2): self.graph[node_id_1][node_id_2]["weight"] = float(weight) self.graph[node_id_1][node_id_2]["distance"] = (None if distance is None else float(distance)) else: self.graph.add_edge( node_id_1, node_id_2, weight=float(weight), distance=(None if distance is None else float(distance)), ) def _to_json_safe(self, x): """Make NetworkX attrs JSON-safe (handles numpy, tuples, ndarrays).""" import numpy as _np if isinstance(x, (int, float, str)) or x is None: return x if isinstance(x, (list, tuple)): return [self._to_json_safe(v) for v in x] if isinstance(x, dict): return {str(k): self._to_json_safe(v) for k, v in x.items()} if isinstance(x, _np.generic): # e.g., np.int64, np.float32 return x.item() if isinstance(x, _np.ndarray): return x.tolist() return str(x) # last-resort fallback def _json_sanitize(self, obj): """ Recursively convert NumPy scalars/arrays, tuples, sets, etc. into JSON-serializable Python types. """ import numpy as np # NumPy scalars -> Python scalars if isinstance(obj, (np.integer,)): return int(obj) if isinstance(obj, (np.floating,)): return float(obj) if isinstance(obj, (np.bool_,)): return bool(obj) # NumPy arrays -> lists if isinstance(obj, np.ndarray): return [self._json_sanitize(x) for x in obj.tolist()] # Containers if isinstance(obj, (list, tuple, set)): return [self._json_sanitize(x) for x in obj] if isinstance(obj, dict): return {str(k): self._json_sanitize(v) for k, v in obj.items()} # Leave JSON-friendly primitives (str, int, float, bool, None) as-is return obj def save_to_json(self, path): """ Save graph to JSON, ensuring every edge has `weight` and `distance`, and all attributes are JSON-serializable. """ import json # Make sure edges have metrics (fills in `distance`, upgrades `weight` if 1) if hasattr(self, "_ensure_edge_metrics"): self._ensure_edge_metrics(overwrite_weight_if_one=True) data = {"nodes": [], "edges": []} # ---- Nodes ---- for n, d in self.graph.nodes(data=True): # prefer unified keys but keep all attrs (sanitized) node_entry = { "id": n, "type": d.get("type") or d.get("node_type"), "position": d.get("position"), "floor": d.get("floor") or d.get("floor_id"), } # include remaining attributes for k, v in d.items(): if k not in node_entry: node_entry[k] = v data["nodes"].append(self._json_sanitize(node_entry)) # ---- Edges ---- for u, v, ed in self.graph.edges(data=True): # build base edge payload edge_entry = { "source": u, "target": v, "weight": ed.get("weight"), "distance": ed.get("distance"), } # include any additional edge attrs for k, v_attr in ed.items(): if k not in edge_entry: edge_entry[k] = v_attr data["edges"].append(self._json_sanitize(edge_entry)) # Write JSON with open(path, "w") as f: json.dump(self._json_sanitize(data), f, indent=2) def plot_on_image( self, image_path, output_path, display_labels=True, threshold_radius=20, highlight_regions=False, ): """ Plot the graph on the given image and save it. Args: image_path (str): Path to the input image. output_path (str): Path to save the plotted image. display_labels (bool): If True, display text labels for nodes. highlight_regions (bool): If True, draw semi-transparent regions around nodes. """ # Load the image image = cv2.imread(image_path) if image is None: raise FileNotFoundError(f"Image not found: {image_path}") overlay = image.copy() # Create an overlay for transparency # BGR colors colors = { # node fills "room": (255, 128, 0), # Bright Orange "door": (0, 204, 102), # Emerald Green "corridor": (255, 102, 255), # Magenta "outside": (204, 51, 51), # Crimson Red "transition": (0, 0, 255), # RED for stairs/elevator (no differentiation) "unknown": (128, 128, 128), # Gray # edge strokes "room_edge": (255, 165, 0), # Lighter Orange "corridor_edge": (51, 153, 255), # Medium Blue "outside_edge": (255, 0, 0), # Bright Red "transition_edge": (0, 0, 180), # Deep RED for any edge touching transition } def _node_color(node_id: str, node_type: str): """Pick a color for the node; stairs/elevator both RED.""" t = (node_type or "").lower() if t in ("transition", "tranistion"): return colors["transition"] return colors.get(t, colors["unknown"]) def _edge_color(type_u: str, type_v: str): """Decide edge color; any transition involvement -> deep RED.""" u = (type_u or "unknown").lower() v = (type_v or "unknown").lower() if "outside" in (u, v): return colors["outside_edge"] if "corridor" in (u, v): return colors["corridor_edge"] if ("transition" in (u, v)) or ("tranistion" in (u, v)): return colors["transition_edge"] return colors["room_edge"] # Step 1: Highlight regions if enabled if highlight_regions: for node_id, data in self.graph.nodes(data=True): if "position" not in data: continue x, y = data["position"] node_type = data.get("type", "unknown") if node_type in {"room", "door", "corridor", "outside", "transition", "tranistion"}: highlight_color = _node_color(node_id, node_type) cv2.circle(overlay, (int(x), int(y)), threshold_radius, highlight_color, -1) # Blend the overlay with the original image for transparency alpha = 0.3 cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0, image) # Step 2: Plot nodes for node_id, data in self.graph.nodes(data=True): if "position" not in data: continue x, y = data["position"] node_type = data.get("type", "unknown") color = _node_color(node_id, node_type) # Size tweaks t = (node_type or "").lower() if t == "corridor": radius = 4 elif t in ("transition", "tranistion"): radius = 9 # a touch larger for visibility else: radius = 8 cv2.circle(image, (int(x), int(y)), radius, color, -1) if display_labels and t != "corridor": cv2.putText( image, str(node_id), (int(x) + 10, int(y) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1, cv2.LINE_AA, ) # Step 3: Plot edges for u, v in self.graph.edges(): if not self.graph.has_node(u) or not self.graph.has_node(v): continue if "position" not in self.graph.nodes[u] or "position" not in self.graph.nodes[v]: continue pos_u = self.graph.nodes[u]["position"] pos_v = self.graph.nodes[v]["position"] type_u = self.graph.nodes[u].get("type", "unknown") type_v = self.graph.nodes[v].get("type", "unknown") edge_color = _edge_color(type_u, type_v) cv2.line( image, (int(pos_u[0]), int(pos_u[1])), (int(pos_v[0]), int(pos_v[1])), edge_color, 2, ) # Save the image cv2.imwrite(output_path, image) print(f"Graph plotted and saved to {output_path}") def add_door_nodes( self, exit_dbboxes, corridor2corridor_dbboxes, room2corridor_dbboxes, room2room_dbboxes, floor_id=None, ): """ Adds door nodes to the graph based on bounding boxes, with unique IDs reflecting the door type. Args: exit_dbboxes (list): List of bounding boxes for exit doors. corridor2corridor_dbboxes (list): List of bounding boxes for corridor-to-corridor doors. room2corridor_dbboxes (list): List of bounding boxes for room-to-corridor doors. room2room_dbboxes (list): List of bounding boxes for room-to-room doors. floor_id (str, optional): Floor id to assign to all created door nodes (if provided). """ # Define door types and their corresponding bounding boxes door_types = [ ("exit", exit_dbboxes), ("c2c", corridor2corridor_dbboxes), ("r2c", room2corridor_dbboxes), ("r2r", room2room_dbboxes), ] # Initialize counters for each door type node_counters = {door_type: 1 for door_type, _ in door_types} for door_type, dbboxes in door_types: centers = [] for bbox in dbboxes: x_center = (bbox[0] + bbox[2]) // 2 y_center = (bbox[1] + bbox[3]) // 2 centers.append((x_center, y_center)) for x, y in centers: node_id = f"{door_type}_door_{node_counters[door_type]}" self.add_node(node_id, "door", (x, y), floor_id=floor_id) node_counters[door_type] += 1 print(f"{len(centers)}-{door_type} door nodes added!") def make_room_door_edges(self, image_path, bboxes): """ Associate door bboxes to MAIN rooms via flood-overlap, then for each MAIN room choose ONE anchor door (closest to the main room). Create exactly ONE edge from that door to the CLOSEST node in the whole family (main or subnode). Do NOT create direct edges to all subnodes; the rest will connect via shortest paths later. Returns: dict: Mapping from bbox tuple -> list of associated MAIN room ids. """ import math if not os.path.exists(image_path): raise FileNotFoundError(f"Image not found: {image_path}") image = cv2.imread(image_path) if image is None: raise ValueError(f"Unable to load image: {image_path}") H, W = image.shape[:2] floodfilled_image = image.copy() # ---- helpers ---- def _is_room(d): return d.get("type") == "room" def _is_sub(n, d): return _is_room(d) and (d.get("is_subnode", False) or "_subnode_" in str(n)) def _parent(n, d): if not _is_sub(n, d): return None p = d.get("parent_room_id") if p: return p s = str(n) return s.split("_subnode_")[0] if "_subnode_" in s else None def _pos(n): return self.graph.nodes[n].get("position") # MAIN rooms only for flood association main_rooms = [(nid, d["position"]) for nid, d in self.graph.nodes(data=True) if _is_room(d) and not _is_sub(nid, d) and "position" in d] door_nodes = [(nid, d["position"]) for nid, d in self.graph.nodes(data=True) if d.get("type") == "door" and "position" in d] if not door_nodes: print("No door nodes present; skipping room↔door association.") return {} door_center_to_id = {(int(px), int(py)): did for did, (px, py) in door_nodes} # ---- flood per main room to get interior pixels ---- flooded_pixels = {} # room_id -> set[(x,y)] point_step = 90 seed_r = 10 for rid, (x, y) in main_rooms: x, y = int(x), int(y) x = max(0, min(x, W - 1)); y = max(0, min(y, H - 1)) visited = set() seed_pts = [(x, y)] for ang in range(0, 360, point_step): rad = np.radians(ang) sx = int(x + seed_r * np.cos(rad)) sy = int(y + seed_r * np.sin(rad)) if 0 <= sx < W and 0 <= sy < H: seed_pts.append((sx, sy)) for sx, sy in seed_pts: if (sx, sy) in visited: continue mask = np.zeros((H + 2, W + 2), np.uint8) _, _, _, rect = cv2.floodFill( floodfilled_image, mask, (sx, sy), (0, 0, 255), loDiff=(10,10,10), upDiff=(10,10,10) ) y0, x0 = max(rect[1], 0), max(rect[0], 0) y1, x1 = min(rect[1]+rect[3], H), min(rect[0]+rect[2], W) for py in range(y0, y1): for px in range(x0, x1): if mask[py+1, px+1] != 0: visited.add((px, py)) flooded_pixels[rid] = visited # ---- associate bboxes to rooms ---- bbox_to_room = {} room_to_doors = {rid: set() for rid, _ in main_rooms} def _doors_in_bbox(x1, y1, x2, y2): cx, cy = (x1 + x2)//2, (y1 + y2)//2 did = door_center_to_id.get((int(cx), int(cy))) if did: return [did] found = [] for did2, (dx, dy) in door_nodes: if x1 <= int(dx) <= x2 and y1 <= int(dy) <= y2: found.append(did2) return found for (x1, y1, x2, y2) in bboxes: x1c, y1c = max(0, x1), max(0, y1) x2c, y2c = min(W - 1, x2), min(H - 1, y2) if x2c < x1c or y2c < y1c: continue bbox_pixels = {(x, y) for x in range(x1c, x2c+1) for y in range(y1c, y2c+1)} associated = [] for rid, pixset in flooded_pixels.items(): if pixset & bbox_pixels: associated.append(rid) if not associated: continue bbox_to_room[(x1, y1, x2, y2)] = associated dids = _doors_in_bbox(x1c, y1c, x2c, y2c) if not dids: print(f"Warning: door bbox {(x1,y1,x2,y2)} matched rooms {associated} but no door node found.") continue for rid in associated: for did in dids: room_to_doors[rid].add(did) # ---- for each room: pick ONE r2c anchor door, connect to ALL r2r doors ---- for rid, _ in main_rooms: dids = list(room_to_doors.get(rid, [])) if not dids: continue # Separate doors by type (check node ID prefix) r2c_doors = [did for did in dids if str(did).startswith("r2c_door_")] r2r_doors = [did for did in dids if str(did).startswith("r2r_door_")] exit_doors = [did for did in dids if str(did).startswith("exit_door_")] # Get family nodes (main + all subnodes) family = [rid] + [n for n, d in self.graph.nodes(data=True) if _is_sub(n, d) and _parent(n, d) == rid and "position" in d] # (1) Pick the nearest r2c door as anchor (exactly ONE) # Connect r2c door to the CLOSEST family node (main or subnode) # The funneling will ensure all subnodes go through main room to reach the door anchor_door = None if r2c_doors: rx, ry = _pos(rid) best_door, best_d = None, float("inf") for did in r2c_doors: dx, dy = _pos(did) d = math.hypot(rx - dx, ry - dy) if d < best_d: best_d, best_door = d, did anchor_door = best_door # Connect r2c door to the CLOSEST family node (optimal connection point) # Funneling will create paths from all subnodes through main room to reach the door if anchor_door is not None: dx, dy = _pos(anchor_door) nearest_node, nearest_dist = None, float("inf") for nid in family: sx, sy = _pos(nid) d = math.hypot(sx - dx, sy - dy) if d < nearest_dist: nearest_dist, nearest_node = d, nid if nearest_node is not None: # Remove any existing edges from other family nodes to this door for other_node in family: if other_node != nearest_node and self.graph.has_edge(other_node, anchor_door): self.graph.remove_edge(other_node, anchor_door) # Connect door to closest family node if not self.graph.has_edge(nearest_node, anchor_door): self.graph.add_edge(nearest_node, anchor_door, weight=float(nearest_dist)) # Store which family node is closest to this door (for reference in funneling) self.graph.nodes[anchor_door]["closest_family_node"] = nearest_node # (2) Connect to ALL r2r doors associated with this room for r2r_door in r2r_doors: dx, dy = _pos(r2r_door) nearest_node, nearest_dist = None, float("inf") for nid in family: sx, sy = _pos(nid) d = math.hypot(sx - dx, sy - dy) if d < nearest_dist: nearest_dist, nearest_node = d, nid if nearest_node is not None and not self.graph.has_edge(nearest_node, r2r_door): self.graph.add_edge(nearest_node, r2r_door, weight=float(nearest_dist)) # (3) Also handle exit doors (connect to closest family node) for exit_door in exit_doors: dx, dy = _pos(exit_door) nearest_node, nearest_dist = None, float("inf") for nid in family: sx, sy = _pos(nid) d = math.hypot(sx - dx, sy - dy) if d < nearest_dist: nearest_dist, nearest_node = d, nid if nearest_node is not None and not self.graph.has_edge(nearest_node, exit_door): self.graph.add_edge(nearest_node, exit_door, weight=float(nearest_dist)) # Store anchor for downstream (r2c door if available, otherwise None) self.graph.nodes[rid]["anchor_door"] = anchor_door print("Doors associated: each room has exactly one r2c door edge, plus all associated r2r door edges.") return bbox_to_room def add_corridor_nodes(self, image_path, corridor_pixels, test_img_dir, dest="corridor", distance=20): """ Processes an image to overlay corridor pixels, create a wall mask, buffer the wall mask, identify invalid pixels, refine the corridor pixel list, and select pixels based on a grid step. Adds selected pixels to the graph and constructs grid-style edges with cross-diagonals. Args: image_path (str): Path to the input image. corridor_pixels (list): List of (y, x) coordinates representing corridor pixels. test_img_dir (str): Directory to save the output images. distance (int): Minimum distance between selected pixels (grid step size). Returns: list[(y,x)]: selected pixels """ # Verify the image exists if not os.path.exists(image_path): raise FileNotFoundError(f"Image not found: {image_path}") # Load the image image = cv2.imread(image_path) if image is None: raise ValueError(f"Unable to load image: {image_path}") # Create the test image directory if it doesn't exist os.makedirs(test_img_dir, exist_ok=True) # 1. Plot the corridor pixels on the image (corridor pixels are in (y, x) format) corridor_overlay_image = image.copy() for y, x in corridor_pixels: cv2.circle(corridor_overlay_image, (x, y), 1, (0, 255, 0), -1) # Green dots for corridor pixels if dest == "corridor": corridor_overlay_path = os.path.join(test_img_dir, "corridor_pixel_overlay.png") else: corridor_overlay_path = os.path.join(test_img_dir, "outside_pixel_overlay.png") cv2.imwrite(corridor_overlay_path, corridor_overlay_image) # 2. Threshold the input image at 240 to create a binary wall mask gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) _, wall_mask = cv2.threshold(gray_image, 240, 255, cv2.THRESH_BINARY) # Invert the wall mask inverted_wall_mask = cv2.bitwise_not(wall_mask) if dest == "corridor": inverted_wall_mask_path = os.path.join(test_img_dir, "corridor_inverted_wall_mask.png") else: inverted_wall_mask_path = os.path.join(test_img_dir, "outside_inverted_wall_mask.png") cv2.imwrite(inverted_wall_mask_path, inverted_wall_mask) # 3. Buffer the inverted wall mask buffered_wall_mask = cv2.dilate(inverted_wall_mask, np.ones((15, 15), np.uint8)) buffered_wall_mask = cv2.bitwise_not(buffered_wall_mask) if dest == "corridor": buffered_wall_mask_path = os.path.join(test_img_dir, "corridor_buffered_wall_mask.png") else: buffered_wall_mask_path = os.path.join(test_img_dir, "outside_buffered_wall_mask.png") cv2.imwrite(buffered_wall_mask_path, buffered_wall_mask) buffered_wall_mask = cv2.bitwise_not(buffered_wall_mask) # 4. Identify invalid pixels and refine the corridor pixel list buffered_wall_coords = set(zip(*np.where(buffered_wall_mask == 255))) # Get (y, x) of wall pixels corridor_set = set(map(tuple, corridor_pixels)) # Convert each [y, x] to (y, x) before creating a set invalid_pixels = corridor_set & buffered_wall_coords # Intersection of wall pixels and corridor pixels refined_corridor_pixels = list(corridor_set - invalid_pixels) # Remove invalids from corridor pixels refined_set = set(refined_corridor_pixels) # Convert to set for fast lookup selected_pixels = [] # Generate grid points and filter them for y in range(0, image.shape[0], distance): for x in range(0, image.shape[1], distance): if (y, x) in refined_set: # Keep only points that exist in refined_corridor_pixels selected_pixels.append((y, x)) # 8. Save refined corridor pixels image refined_corridor_image = image.copy() for y, x in invalid_pixels: cv2.circle(refined_corridor_image, (x, y), 1, (0, 0, 255), -1) # Red for invalid pixels for y, x in refined_corridor_pixels: cv2.circle(refined_corridor_image, (x, y), 1, (0, 255, 0), -1) # Green for valid pixels if dest == "corridor": refined_corridor_path = os.path.join(test_img_dir, "refined_corridor_pixels.png") else: refined_corridor_path = os.path.join(test_img_dir, "refined_outside_pixels.png") cv2.imwrite(refined_corridor_path, refined_corridor_image) # 9. Save selected pixel map image selected_pixel_image = image.copy() for y, x in selected_pixels: cv2.circle(selected_pixel_image, (x, y), 4, (139, 0, 139), -1) # Dark blue for selected pixels if dest == "corridor": selected_pixel_map_path = os.path.join(test_img_dir, "selected_corridor_pixel_map.png") else: selected_pixel_map_path = os.path.join(test_img_dir, "selected_outside_pixel_map.png") cv2.imwrite(selected_pixel_map_path, selected_pixel_image) return selected_pixels def add_corridor_edges(self, selected_pixels, distance=20): """ Adds corridor edges based on the selected pixels and grid distance. Args: selected_pixels (list): List of (y, x) coordinates representing corridor pixels. distance (int): The grid step size for connecting nodes. """ # Step 2: Add edges between corridor nodes selected_pixel_positions = {(y, x): f"corridor_connect_{i + 1}" for i, (y, x) in enumerate(selected_pixels)} for y, x in selected_pixel_positions.keys(): node_id = selected_pixel_positions[(y, x)] # Define neighbor offsets (horizontal, vertical, and diagonal) neighbors = [ (y + distance, x), # Down (y - distance, x), # Up (y, x + distance), # Right (y, x - distance), # Left (y + distance, x + distance), # Bottom-right diagonal (y - distance, x - distance), # Top-left diagonal (y + distance, x - distance), # Bottom-left diagonal (y - distance, x + distance), # Top-right diagonal ] # Add edges if the neighbor exists in the grid AND both nodes already exist for ny, nx in neighbors: if (ny, nx) in selected_pixel_positions: neighbor_id = selected_pixel_positions[(ny, nx)] if self.graph.has_node(node_id) and self.graph.has_node(neighbor_id): self.graph.add_edge(node_id, neighbor_id) def add_outdoor_edges(self, outdoor_pixels, distance=20): """ Adds outdoor edges based on the selected pixels and grid distance. Args: outdoor_pixels (list): List of (y, x) coordinates representing outdoor pixels. distance (int): The grid step size for connecting nodes. """ selected_pixel_positions = {(y, x): f"outside_connect_{i + 1}" for i, (y, x) in enumerate(outdoor_pixels)} for y, x in selected_pixel_positions.keys(): node_id = selected_pixel_positions[(y, x)] neighbors = [ (y + distance, x), # Down (y - distance, x), # Up (y, x + distance), # Right (y, x - distance), # Left (y + distance, x + distance), # Bottom-right diagonal (y - distance, x - distance), # Top-left diagonal (y + distance, x - distance), # Bottom-left diagonal (y - distance, x + distance), # Top-right diagonal ] for ny, nx in neighbors: if (ny, nx) in selected_pixel_positions: neighbor_id = selected_pixel_positions[(ny, nx)] if self.graph.has_node(node_id) and self.graph.has_node(neighbor_id): self.graph.add_edge(node_id, neighbor_id) def connect_hallways(self): print("Connecting hallways...") # Find all corridor main nodes corridor_main_nodes = [ node for node, data in self.graph.nodes(data=True) if data.get('type') == 'corridor' and str(node).startswith("corridor_main") ] # Find all corridor connect nodes corridor_connect_nodes = { node for node, data in self.graph.nodes(data=True) if data.get('type') == 'corridor' and str(node).startswith("corridor_connect") } # Find all outside main nodes outside_main_nodes = [ node for node, data in self.graph.nodes(data=True) if data.get('type') == 'outside' and str(node).startswith("outside_main") ] # Find all outside connect nodes outside_connect_nodes = { node for node, data in self.graph.nodes(data=True) if data.get('type') == 'outside' and str(node).startswith("outside_connect") } # Euclidean distance on node positions def euclidean_distance(node_1, node_2): pos_1 = self.graph.nodes[node_1].get('position', [0, 0]) pos_2 = self.graph.nodes[node_2].get('position', [0, 0]) return math.sqrt((pos_1[0] - pos_2[0]) ** 2 + (pos_1[1] - pos_2[1]) ** 2) # Radius within which we look for corridor_connect or outside_connect nodes radius = 200 # For each main corridor node, connect to up to 4 nearby corridor_connect nodes for main_node in corridor_main_nodes: nearby = [] for connect_node in corridor_connect_nodes: dist = euclidean_distance(main_node, connect_node) if dist <= radius: nearby.append((dist, connect_node)) if nearby: for dist, connect_node in sorted(nearby, key=lambda x: x[0])[:4]: if not self.graph.has_edge(main_node, connect_node): self.add_edge(main_node, connect_node, weight=dist) else: print(f"No corridor_connect nodes found within radius of {main_node}") # For each main outside node, connect to up to 4 nearby outside_connect nodes for main_node in outside_main_nodes: nearby = [] for connect_node in outside_connect_nodes: dist = euclidean_distance(main_node, connect_node) if dist <= radius: nearby.append((dist, connect_node)) if nearby: for dist, connect_node in sorted(nearby, key=lambda x: x[0])[:4]: if not self.graph.has_edge(main_node, connect_node): self.add_edge(main_node, connect_node, weight=dist) else: print(f"No outside_connect nodes found within radius of {main_node}") print(f"Added edges to {len(corridor_main_nodes)} corridor main nodes.") print(f"Added edges to {len(outside_main_nodes)} outside main nodes.") def connect_doors(self): print("\nConnecting doors...") # Door nodes by category exit_doors = [ node for node, data in self.graph.nodes(data=True) if data.get('type') == 'door' and str(node).startswith("exit_door") ] c2c_doors = [ node for node, data in self.graph.nodes(data=True) if data.get('type') == 'door' and str(node).startswith("c2c_door") ] r2c_doors = [ node for node, data in self.graph.nodes(data=True) if data.get('type') == 'door' and str(node).startswith("r2c_door") ] # Corridor/outside connectivity targets outside_connect_nodes = { node for node, data in self.graph.nodes(data=True) if data.get('type') == 'outside' and str(node).startswith("outside_connect") } corridor_connect_nodes = { node for node, data in self.graph.nodes(data=True) if data.get('type') == 'corridor' and str(node).startswith("corridor_connect") } corridor_main_nodes = { node for node, data in self.graph.nodes(data=True) if data.get('type') == 'corridor' and str(node).startswith("corridor_main") } def euclidean_distance(node_1, node_2): pos_1 = self.graph.nodes[node_1].get('position', [0, 0]) pos_2 = self.graph.nodes[node_2].get('position', [0, 0]) return math.sqrt((pos_1[0] - pos_2[0]) ** 2 + (pos_1[1] - pos_2[1]) ** 2) radius = 100 # Exit doors: connect to nearest outside_connect and corridor_connect for door in exit_doors: nearby_out = [] nearby_cor = [] for connect_node in outside_connect_nodes: dist = euclidean_distance(door, connect_node) if dist <= radius: nearby_out.append((dist, connect_node)) for connect_node in corridor_connect_nodes: dist = euclidean_distance(door, connect_node) if dist <= radius: nearby_cor.append((dist, connect_node)) if nearby_out: dist, cn = min(nearby_out, key=lambda x: x[0]) if not self.graph.has_edge(door, cn): self.add_edge(door, cn, weight=dist) else: print(f"No outside_connect nodes found within radius of {door}") if nearby_cor: dist, cn = min(nearby_cor, key=lambda x: x[0]) if not self.graph.has_edge(door, cn): self.add_edge(door, cn, weight=dist) else: print(f"No corridor_connect nodes found within radius of {door}") # c2c doors: connect to up to 4 nearest corridor_connect nodes for door in c2c_doors: nearby = [] for connect_node in corridor_connect_nodes: dist = euclidean_distance(door, connect_node) if dist <= radius: nearby.append((dist, connect_node)) if nearby: for dist, cn in sorted(nearby, key=lambda x: x[0])[:4]: if not self.graph.has_edge(door, cn): self.add_edge(door, cn, weight=dist) else: print(f"No corridor_connect nodes found within radius of {door}") # r2c doors: connect to nearest of corridor_connect or corridor_main for door in r2c_doors: nearby = [] for cn in corridor_connect_nodes: dist = euclidean_distance(door, cn) if dist <= radius: nearby.append((dist, cn)) for mn in corridor_main_nodes: dist = euclidean_distance(door, mn) if dist <= radius: nearby.append((dist, mn)) if nearby: dist, tgt = min(nearby, key=lambda x: x[0]) if not self.graph.has_edge(door, tgt): self.add_edge(door, tgt, weight=dist) else: print(f"No corridor_connect or corridor_main nodes found within radius of {door}") def connect_rooms(self): print("\nConnecting rooms...") # Helper: is this a room subnode? def _is_room_subnode(node_id, data): return data.get('type') == 'room' and (data.get('is_subnode', False) or "_subnode_" in str(node_id)) # Only MAIN room nodes (exclude densified subnodes) room_nodes = [ node for node, data in self.graph.nodes(data=True) if data.get('type') == 'room' and not _is_room_subnode(node, data) ] corridor_main_nodes = { node for node, data in self.graph.nodes(data=True) if data.get('type') == 'corridor' and str(node).startswith("corridor_main") } corridor_connect_nodes = { node for node, data in self.graph.nodes(data=True) if data.get('type') == 'corridor' and str(node).startswith("corridor_connect") } def euclidean_distance(node_1, node_2): pos_1 = self.graph.nodes[node_1].get('position', [0, 0]) pos_2 = self.graph.nodes[node_2].get('position', [0, 0]) return math.sqrt((pos_1[0] - pos_2[0]) ** 2 + (pos_1[1] - pos_2[1]) ** 2) radius = 400 disconnected_room_count = 0 for room in room_nodes: # If this main room has no neighbors, try to tie it to the corridor graph if len(list(self.graph.neighbors(room))) == 0: disconnected_room_count += 1 nearby = [] for cn in corridor_main_nodes: dist = euclidean_distance(room, cn) if dist <= radius: nearby.append((dist, cn)) for cn in corridor_connect_nodes: dist = euclidean_distance(room, cn) if dist <= radius: nearby.append((dist, cn)) if nearby: dist, closest = min(nearby, key=lambda x: x[0]) if not self.graph.has_edge(room, closest): self.add_edge(room, closest, weight=dist) else: print(f"No corridor nodes found within radius of {room}") print(f"Total disconnected main room nodes: {disconnected_room_count}") def connect_transitions(self): """ Connect transition nodes (stairs/elevators) to the graph by connecting them to nearby corridors or rooms. This ensures transitions appear in pre-pruning plots. """ print("\nConnecting transitions...") # Get all transition nodes transition_nodes = [ node for node, data in self.graph.nodes(data=True) if data.get('type') == 'transition' ] if not transition_nodes: print("No transition nodes found.") return # Get corridor nodes (both main and connect types) corridor_main_nodes = { node for node, data in self.graph.nodes(data=True) if data.get('type') == 'corridor' and str(node).startswith("corridor_main") } corridor_connect_nodes = { node for node, data in self.graph.nodes(data=True) if data.get('type') == 'corridor' and str(node).startswith("corridor_connect") } def euclidean_distance(node_1, node_2): pos_1 = self.graph.nodes[node_1].get('position', [0, 0]) pos_2 = self.graph.nodes[node_2].get('position', [0, 0]) return math.sqrt((pos_1[0] - pos_2[0]) ** 2 + (pos_1[1] - pos_2[1]) ** 2) radius = 400 connected_count = 0 for transition in transition_nodes: # Check if already connected if len(list(self.graph.neighbors(transition))) > 0: continue # Find nearby corridors nearby = [] for cn in corridor_main_nodes: dist = euclidean_distance(transition, cn) if dist <= radius: nearby.append((dist, cn)) for cn in corridor_connect_nodes: dist = euclidean_distance(transition, cn) if dist <= radius: nearby.append((dist, cn)) if nearby: dist, closest = min(nearby, key=lambda x: x[0]) if not self.graph.has_edge(transition, closest): self.add_edge(transition, closest, weight=dist) connected_count += 1 else: print(f"No corridor nodes found within radius of {transition}") print(f"Connected {connected_count} transition nodes to the graph.") def merge_nearby_nodes(self, threshold_room=50, threshold_door=30): """ Merge nodes that are within a certain vicinity threshold, with different thresholds for "room" and "door" nodes. """ # Build a list of node IDs and their positions node_ids = list(self.graph.nodes) positions = {node_id: self.graph.nodes[node_id]["position"] for node_id in node_ids} # Separate nodes by type to apply different thresholds room_nodes = [node_id for node_id in node_ids if self.graph.nodes[node_id].get("type") == "room"] door_nodes = [node_id for node_id in node_ids if self.graph.nodes[node_id].get("type") == "door"] # Initialize Union-Find structure for nodes parent = {node_id: node_id for node_id in room_nodes + door_nodes} def find(u): if parent[u] != u: parent[u] = find(parent[u]) return parent[u] def union(u, v): pu, pv = find(u), find(v) if pu != pv: parent[pv] = pu # Merging room nodes for i in range(len(room_nodes)): node_id_1 = room_nodes[i] pos_1 = np.array(positions[node_id_1]) for j in range(i + 1, len(room_nodes)): node_id_2 = room_nodes[j] pos_2 = np.array(positions[node_id_2]) dist = np.linalg.norm(pos_1 - pos_2) if dist < threshold_room: union(node_id_1, node_id_2) # Merging door nodes for i in range(len(door_nodes)): node_id_1 = door_nodes[i] pos_1 = np.array(positions[node_id_1]) for j in range(i + 1, len(door_nodes)): node_id_2 = door_nodes[j] pos_2 = np.array(positions[node_id_2]) dist = np.linalg.norm(pos_1 - pos_2) if dist < threshold_door: union(node_id_1, node_id_2) # Group nodes by their representative parent node clusters = {} for node_id in room_nodes + door_nodes: p = find(node_id) clusters.setdefault(p, []).append(node_id) # Merge nodes in each cluster for cluster_nodes in clusters.values(): if len(cluster_nodes) > 1: positions_list = [positions[node_id] for node_id in cluster_nodes] avg_position = tuple(map(int, np.mean(positions_list, axis=0))) # Keep the first node as the main node main_node = cluster_nodes[0] self.graph.nodes[main_node]["position"] = avg_position # Remove other nodes from the graph and node_types for node_id in cluster_nodes[1:]: self.graph.remove_node(node_id) for node_type, node_list in self.node_types.items(): if node_id in node_list: node_list.remove(node_id) print(f"Merged nodes {cluster_nodes} into {main_node} at {avg_position}") return self.graph def _ensure_edge_weights(self): """ Ensure every edge has a numeric 'weight' for shortest-path queries. If missing, use Euclidean distance between node positions; fall back to 1.0. """ import math for u, v, data in self.graph.edges(data=True): if "weight" not in data or data["weight"] is None: pu = self.graph.nodes[u].get("position") pv = self.graph.nodes[v].get("position") if pu is not None and pv is not None: w = math.hypot(float(pu[0]) - float(pv[0]), float(pu[1]) - float(pv[1])) else: w = 1.0 data["weight"] = float(w) def connect_all_rooms(self, input_path, graph_img_dir): """ PRUNING with guarantees: 1) For each room family (main + subnodes): - Ensure an anchor exists: use main's 'anchor_door' if present; otherwise add ONE edge from the nearest corridor (same floor preferred) to the closest family node. - Keep the union of SHORTEST PATHS (legal, weighted) from every family node (main + subs) to that anchor. Only INTRA-FAMILY edges on those paths are kept; corridor/door segments used by those paths are kept too. 2) Global connectivity/pruning: - Keep the union of SHORTEST PATHS between every pair of MAIN room nodes. - For each EXIT DOOR, keep the shortest path from the closest MAIN room. - NEW: For each TRANSITION node (stairs/elevator), keep the union of SHORTEST PATHS from EVERY MAIN room to that transition node (all pairs main→transition). 3) Prune: remove all nodes/edges not on any kept path. 4) NEW: Transition nodes are never pruned. After pruning, ensure each transition is attached to its nearest remaining corridor (fallback connection), if not already. """ import math, random import numpy as np import networkx as nx from PIL import Image import matplotlib.pyplot as plt # ---------- helpers ---------- def _is_room(d): return d.get('type') == 'room' def _is_sub(n, d): return _is_room(d) and (d.get('is_subnode', False) or "_subnode_" in str(n)) def _parent(n, d): if not _is_sub(n, d): return None p = d.get("parent_room_id") if p: return p s = str(n) return s.split("_subnode_")[0] if "_subnode_" in s else None def _pos(n): return self.graph.nodes[n].get('position') if n in self.graph else None def _dist(p, q): if p is None or q is None: return float('inf') return ((p[0]-q[0])**2 + (p[1]-q[1])**2) ** 0.5 # Make sure weighted shortest paths reflect Euclidean-ish lengths self._ensure_edge_weights() # ---------- collect entities ---------- main_rooms = [n for n, d in self.graph.nodes(data=True) if _is_room(d) and not _is_sub(n, d)] subrooms = [n for n, d in self.graph.nodes(data=True) if _is_sub(n, d)] corridors = [n for n, d in self.graph.nodes(data=True) if d.get('type') == 'corridor'] exit_doors = [n for n, d in self.graph.nodes(data=True) if d.get('type') == 'door' and str(n).startswith("exit_door")] # NEW: strictly 'transition' (no legacy 'tranistion') transitions = [n for n, d in self.graph.nodes(data=True) if d.get('type') == 'transition'] room_family = {rid: [rid] for rid in main_rooms} for n in subrooms: pr = _parent(n, self.graph.nodes[n]) if pr in room_family: room_family[pr].append(n) print("\nConnecting all rooms (with pruning to shortest paths)...") print( f"Totals -> nodes: {len(self.graph.nodes)}, " f"main rooms: {len(main_rooms)}, subrooms: {len(subrooms)}, " f"corridors: {len(corridors)}, exits: {len(exit_doors)}, transitions: {len(transitions)}" ) # ---------- per-family: ensure anchor & build union of shortest-to-anchor ---------- family_keep_nodes, family_keep_edges = set(), set() corridor_fallback_used = 0 path_segments_to_plot = [] for rid, fam in room_family.items(): fam = [n for n in fam if n in self.graph] if not fam: continue # (a) determine/create anchor anchor = self.graph.nodes[rid].get("anchor_door") if anchor is not None and anchor not in self.graph: anchor = None if anchor is None: # no door known -> attach nearest corridor to the closest family node rid_floor = self.graph.nodes[rid].get('floor') same_floor = [c for c in corridors if self.graph.nodes[c].get('floor') == rid_floor] candidates = same_floor if same_floor else corridors if candidates: best_pair, best_d = None, float('inf') for fn in fam: p = _pos(fn) if p is None: continue for cn in candidates: d = _dist(p, _pos(cn)) if d < best_d: best_d, best_pair = d, (fn, cn) if best_pair is not None: fn, cn = best_pair if not self.graph.has_edge(fn, cn): self.graph.add_edge(fn, cn, weight=float(best_d)) anchor = cn corridor_fallback_used += 1 else: continue else: continue # (b) union of shortest paths from every family node to anchor for n in fam: if n == anchor: continue try: sp = nx.shortest_path(self.graph, source=n, target=anchor, weight='weight') except nx.NetworkXNoPath: # micro-fix: stitch nearest family mate then retry once pn = _pos(n) best_mate, best_d = None, float('inf') for m in fam: if m == n: continue d = _dist(pn, _pos(m)) if d < best_d: best_d, best_mate = d, m if best_mate is not None and not self.graph.has_edge(n, best_mate): self.graph.add_edge(n, best_mate, weight=float(best_d)) try: sp = nx.shortest_path(self.graph, source=n, target=anchor, weight='weight') except nx.NetworkXNoPath: print(f"[{rid}] no path from {n} to anchor after local fix; skipping this node.") continue else: print(f"[{rid}] no path from {n} to anchor; skipping this node.") continue family_keep_nodes.update(sp) path_segments_to_plot.append(sp) for u, v in zip(sp[:-1], sp[1:]): family_keep_edges.add((u, v) if u < v else (v, u)) if corridor_fallback_used: print(f"Corridor fallback used for {corridor_fallback_used} families lacking doors.") # ---------- global: shortest paths between EVERY pair of MAIN rooms ---------- global_keep_nodes, global_keep_edges = set(), set() main_pairs_no_path = 0 for i, a in enumerate(main_rooms): for b in main_rooms[i+1:]: try: sp = nx.shortest_path(self.graph, source=a, target=b, weight='weight') global_keep_nodes.update(sp) path_segments_to_plot.append(sp) for u, v in zip(sp[:-1], sp[1:]): global_keep_edges.add((u, v) if u < v else (v, u)) except nx.NetworkXNoPath: main_pairs_no_path += 1 if main_pairs_no_path: print(f"WARNING: {main_pairs_no_path} main-room pairs had no path before pruning (graph may be fragmented).") # ---------- exit doors: keep shortest path from closest MAIN room ---------- exit_keep_nodes, exit_keep_edges = set(), set() for ed in exit_doors: best_sp, best_len = None, float('inf') for rid in main_rooms: try: sp = nx.shortest_path(self.graph, source=rid, target=ed, weight='weight') if len(sp) < best_len: best_len, best_sp = len(sp), sp except nx.NetworkXNoPath: continue if best_sp: exit_keep_nodes.update(best_sp) path_segments_to_plot.append(best_sp) for u, v in zip(best_sp[:-1], best_sp[1:]): exit_keep_edges.add((u, v) if u < v else (v, u)) # ---------- NEW: transitions - keep SHORTEST PATHS from EVERY MAIN room to EVERY transition ---------- transition_keep_nodes, transition_keep_edges = set(), set() trans_pairs_no_path = 0 for t in transitions: for rid in main_rooms: try: sp = nx.shortest_path(self.graph, source=rid, target=t, weight='weight') transition_keep_nodes.update(sp) path_segments_to_plot.append(sp) for u, v in zip(sp[:-1], sp[1:]): transition_keep_edges.add((u, v) if u < v else (v, u)) except nx.NetworkXNoPath: trans_pairs_no_path += 1 if trans_pairs_no_path: print(f"Note: {trans_pairs_no_path} main→transition pairs had no path before pruning.") # ---------- build final KEEP sets & prune ---------- keep_nodes = ( family_keep_nodes | global_keep_nodes | exit_keep_nodes | transition_keep_nodes | set(transitions) # NEVER prune transitions ) keep_edges = family_keep_edges | global_keep_edges | exit_keep_edges | transition_keep_edges # Ensure endpoints of kept edges are kept for u, v in list(keep_edges): keep_nodes.add(u); keep_nodes.add(v) # Remove nodes not in keep nodes_to_remove = set(self.graph.nodes) - keep_nodes if nodes_to_remove: self.graph.remove_nodes_from(nodes_to_remove) # Remove edges not in keep (and re-check endpoints) edges_to_remove = [] for u, v in self.graph.edges(): e = (u, v) if u < v else (v, u) if (u not in keep_nodes) or (v not in keep_nodes) or (e not in keep_edges): edges_to_remove.append((u, v)) if edges_to_remove: self.graph.remove_edges_from(edges_to_remove) print(f"After pruning -> nodes: {len(self.graph.nodes)}, edges: {len(self.graph.edges)}") # ---------- Fallback: attach isolated transitions to nearest remaining corridor ---------- post_corridors = [n for n, d in self.graph.nodes(data=True) if d.get('type') == 'corridor'] attached = 0 if post_corridors: for tn in transitions: if tn not in self.graph: continue # If already connected to any corridor, skip if any(self.graph.nodes[nbr].get('type') == 'corridor' for nbr in self.graph.neighbors(tn)): continue pt = _pos(tn) if pt is None: continue best_c, best_d = None, float('inf') for cn in post_corridors: d = _dist(pt, _pos(cn)) if d < best_d: best_d, best_c = d, cn if best_c is not None and not math.isinf(best_d): self.graph.add_edge(tn, best_c, weight=float(best_d)) path_segments_to_plot.append([tn, best_c]) attached += 1 else: print("WARNING: No corridors remain after pruning; transition nodes were kept but not connected.") if attached: print(f"Transition attachments added post-pruning: {attached}") # ---------- visualize kept paths ---------- try: img = Image.open(input_path) w, h = img.size fig, ax = plt.subplots(figsize=(max(1, w/100), max(1, h/100)), dpi=100) ax.imshow(img) random.seed(42) def _clr(): return (random.random(), random.random(), random.random()) for sp in path_segments_to_plot: coords = [] for n in sp: if n in self.graph.nodes: p = self.graph.nodes[n].get('position') if p is not None: coords.append(p) if len(coords) >= 2: arr = np.array(coords) ax.plot(arr[:, 0], arr[:, 1], color=_clr(), linewidth=2) out_path = f"{graph_img_dir}/colored_paths.png" plt.axis('off'); plt.savefig(out_path, bbox_inches='tight', pad_inches=0); plt.close() print(f"Shortest-path visualization saved: {out_path}") except Exception as e: print(f"Plotting skipped: {e}") return self.graph def remove_edges_not_in_web(self, web_nodes): # Create a list of edges to remove edges_to_remove = [ (u, v) for u, v in self.graph.edges() if u not in web_nodes or v not in web_nodes ] self.graph.remove_edges_from(edges_to_remove) def return_graph_size(self): return len(self.graph.nodes) def connect_room_family_funnel(self, room_id: str, spacing_px: int = 60, door_selector: str = "nearest") -> int: """ Build a local lattice inside a room (main + subnodes), then keep only the intra-room edges that lie on shortest paths to the room's anchor(s). Anchors: - Doors attached to any family member (preferred). - If no doors exist, FALL BACK to the nearest corridor node and use the nearest family node to that corridor as the single anchor. Returns: number of kept intra-room edges for this family. """ import math import networkx as nx if room_id not in self.graph: return 0 # -------- collect family (main + subnodes) -------- def _is_room(d): return d.get("type") == "room" def _is_sub(n, d): return _is_room(d) and (d.get("is_subnode", False) or "_subnode_" in str(n)) family = [room_id] for nid, data in self.graph.nodes(data=True): if nid == room_id: continue if not _is_room(data): continue if _is_sub(nid, data): parent = data.get("parent_room_id") if parent == room_id or (parent is None and str(nid).startswith(f"{room_id}_subnode_")): family.append(nid) # positions pos = {} for nid in family: p = self.graph.nodes[nid].get("position") if p is not None: pos[nid] = (float(p[0]), float(p[1])) family = [nid for nid in family if nid in pos] if len(family) <= 1: return 0 # -------- preferred anchors: DOORS attached to ANY family member -------- ext_to_anchor = [] # list of (external_node_id, anchor_family_node_id) for nid in family: for nbr in self.graph.neighbors(nid): if self.graph.nodes[nbr].get("type") == "door": # For r2c doors, always use main room as anchor (paths must go through main room) # For other doors, use nearest family node door_id = str(nbr) if door_id.startswith("r2c_door_") or door_id.startswith("exit_door_"): # r2c and exit doors: anchor is always the main room anchor = room_id else: # r2r and other doors: use nearest family node to this door dp = self.graph.nodes[nbr].get("position") if dp is None: continue dx, dy = float(dp[0]), float(dp[1]) anchor = min(family, key=lambda n: math.hypot(pos[n][0] - dx, pos[n][1] - dy)) ext_to_anchor.append((nbr, anchor)) # -------- FALLBACK: nearest CORRIDOR node if no doors -------- if not ext_to_anchor: # choose nearest corridor node to the room family's centroid fx = sum(pos[n][0] for n in family) / len(family) fy = sum(pos[n][1] for n in family) / len(family) best_corr = None best_d = float("inf") for nid, data in self.graph.nodes(data=True): if data.get("type") != "corridor": continue cp = data.get("position") if cp is None: continue d = math.hypot(float(cp[0]) - fx, float(cp[1]) - fy) if d < best_d: best_d = d best_corr = nid if best_corr is None: return 0 # nothing to funnel to # anchor is the family node closest to this corridor node cpx, cpy = map(float, self.graph.nodes[best_corr]["position"]) anchor = min(family, key=lambda n: math.hypot(pos[n][0] - cpx, pos[n][1] - cpy)) ext_to_anchor.append((best_corr, anchor)) # ensure direct graph edges from each anchor to its external node (door/corridor) for ext, anc in ext_to_anchor: # add weighted edge if missing ep = self.graph.nodes[ext].get("position") if ep is None: continue w = math.hypot(pos[anc][0] - float(ep[0]), pos[anc][1] - float(ep[1])) if not self.graph.has_edge(anc, ext): self.graph.add_edge(anc, ext, weight=float(w)) # -------- build local lattice (short edges only) -------- import numpy as np r = max(2.0, float(spacing_px) * 1.25) # neighbor radius temp_edges = set() def _eudist(a, b): ax, ay = pos[a]; bx, by = pos[b] return math.hypot(ax - bx, ay - by) # lattice subgraph with only family nodes Gf = nx.Graph() for n in family: Gf.add_node(n) for i in range(len(family)): for j in range(i + 1, len(family)): u, v = family[i], family[j] d = _eudist(u, v) if d <= r: Gf.add_edge(u, v, weight=d) if not self.graph.has_edge(u, v): self.graph.add_edge(u, v, weight=float(d), _temp_family_edge=True, _family_owner=room_id) else: ed = self.graph.edges[u, v] ed.setdefault("_temp_family_edge", True) ed["_family_owner"] = room_id temp_edges.add(tuple(sorted((u, v)))) # -------- compute funnel paths inside the room lattice -------- keep_edges = set() if door_selector == "nearest": # precompute SSSP from each anchor node (inside lattice) packs = [] for _, anchor in ext_to_anchor: if anchor not in Gf: continue dist, paths = nx.single_source_dijkstra(Gf, anchor, weight="weight") packs.append((anchor, dist, paths)) for n in family: best_path = None best_cost = float("inf") for anchor, dist, paths in packs: if n in dist and dist[n] < best_cost: best_cost = dist[n] best_path = paths[n] if best_path and len(best_path) > 1: for u, v in zip(best_path[:-1], best_path[1:]): keep_edges.add(tuple(sorted((u, v)))) else: # union to all anchors for _, anchor in ext_to_anchor: if anchor not in Gf: continue dist, paths = nx.single_source_dijkstra(Gf, anchor, weight="weight") for n in family: if n in paths and len(paths[n]) > 1: for u, v in zip(paths[n][:-1], paths[n][1:]): keep_edges.add(tuple(sorted((u, v)))) # -------- prune temporary lattice edges not used by any path -------- for u, v in list(self.graph.edges()): ed = self.graph.edges[u, v] if ed.get("_temp_family_edge") and ed.get("_family_owner") == room_id: if tuple(sorted((u, v))) not in keep_edges: self.graph.remove_edge(u, v) else: ed.pop("_temp_family_edge", None) ed.pop("_family_owner", None) # -------- For r2c doors: ensure ALL family nodes within radius are connected to main room -------- # This guarantees all subnodes can reach r2c doors through the main room has_r2c_door = any(str(ext).startswith("r2c_door_") for ext, _ in ext_to_anchor) if has_r2c_door and room_id in family: # Ensure main room is connected to all subnodes within lattice radius for nid in family: if nid == room_id: continue if nid not in pos: continue dist = _eudist(room_id, nid) if dist <= r: # Within lattice radius if not self.graph.has_edge(room_id, nid): # Add direct connection to main room to ensure connectivity self.graph.add_edge(room_id, nid, weight=float(dist)) # For r2c doors: keep connection to closest family node (optimal placement) # The funneling uses main room as anchor, so it adds main room -> door edge # We keep BOTH connections: # - closest family node -> door (optimal connection point, as user requested) # - main room -> door (funneling anchor) # This allows optimal door placement while maintaining funneling structure # Note: If closest is a subnode, it will have a direct path to door # Other subnodes will go through main room to reach the door for ext, _ in ext_to_anchor: if str(ext).startswith("r2c_door_"): # Get the closest family node (stored when door was initially connected) closest_node = self.graph.nodes[ext].get("closest_family_node") # Remove connections from subnodes that are NOT the closest one # Keep the connection to closest node (optimal placement) # Keep the connection to main room (funneling anchor) for nid in family: if nid != room_id and nid != closest_node and self.graph.has_edge(nid, ext): self.graph.remove_edge(nid, ext) kept = len([e for e in keep_edges if e in temp_edges]) return kept def connect_all_families_funnel(self, spacing_px: int = 60, door_selector: str = "nearest") -> int: """Run the funnel connector for every MAIN room (excludes subnodes).""" total = 0 for nid, data in self.graph.nodes(data=True): if data.get("type") == "room" and not data.get("is_subnode", False) and "_subnode_" not in str(nid): total += self.connect_room_family_funnel(nid, spacing_px=spacing_px, door_selector=door_selector) return total @staticmethod def calculate_bbox_centers(bboxes): """ Calculate the centers of bounding boxes from the given list of bounding boxes. Args: bboxes (list): List of bounding boxes, where each bounding box is represented as a list of 8 coordinates [x1, y1, x2, y2, x3, y3, x4, y4]. Returns: list: List of (x, y) centers for each bounding box. """ bbox_centers = [] for coordinates in bboxes: points = np.array(coordinates).reshape(4, 2) center_x = np.mean(points[:, 0]) center_y = np.mean(points[:, 1]) bbox_centers.append((center_x, center_y)) return bbox_centers def connect_doors_to_rooms(self): """ Connect doors to their nearest rooms. Rules: - A door can only connect to one room (nearest room based on distance). - A room can have multiple doors. """ door_nodes = [node_id for node_id in self.node_types["door"]] room_nodes = [node_id for node_id in self.node_types["room"]] if not door_nodes or not room_nodes: print("No doors or rooms available to connect.") return print("\nConnecting doors to rooms...") for door_id in door_nodes: door_pos = np.array(self.graph.nodes[door_id]["position"]) nearest_room = None min_distance = float("inf") for room_id in room_nodes: room_pos = np.array(self.graph.nodes[room_id]["position"]) distance = np.linalg.norm(door_pos - room_pos) if distance < min_distance: min_distance = distance nearest_room = room_id if nearest_room: self.add_edge(door_id, nearest_room) print(f"Connected door '{door_id}' to room '{nearest_room}' (distance: {min_distance:.2f})")