Tesseract / Main.py
yansari's picture
feat: graph-editor UX overhaul and headless matplotlib fix
ef4c917
Raw
History Blame Contribute Delete
53.3 kB
import os
import sys
import argparse
from pathlib import Path
import torch
import numpy as np
import cv2
import time
import re
import pandas as pd
from PIL import Image
# Force a headless matplotlib backend BEFORE any module imports pyplot.
# The web pipeline runs in a worker thread where interactive (Tk) backends
# crash with "main thread is not in main loop". Agg is non-interactive and safe.
import matplotlib
matplotlib.use("Agg")
# Add the required paths to the Python path
_PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.join(_PROJECT_ROOT, "Models", "Text_Models"))
sys.path.append(os.path.join(_PROJECT_ROOT, "utils"))
sys.path.append(os.path.join(_PROJECT_ROOT, "Models", "Interpreter"))
sys.path.append(os.path.join(_PROJECT_ROOT, "Models", "Door_Models"))
from text_interpreter import interpret_bboxes, parse_transition_labels
from door_bboxer import*
# Import functions and classes
from text_bboxer import*
from utils.graph import BuildingGraph
from utils.floodfill import*
from utils.floodfill import get_room_subnode_candidates
from utils.connectivity import*
from utils.Improve import*
from utils.compute_time_eval import*
def check_image_exists(image_name, input_images_dir):
"""
Check if the given image exists in the input images directory.
Args:
image_name (str): Name of the image with the extension.
input_images_dir (str): Path to the input images directory.
Returns:
str: Full path to the image if it exists.
Raises:
FileNotFoundError: If the image does not exist.
"""
image_path = Path(input_images_dir) / image_name
if not image_path.exists():
raise FileNotFoundError(f"Error: The image '{image_name}' does not exist in '{input_images_dir}'.")
return str(image_path)
def detect_floor_from_filename(image_name):
"""
Extract floor number from image filename.
Rules:
- First word after splitting by spaces indicates the floor
- "FF" → 1 (First Floor, lowest)
- "SF" → 2 (Second Floor)
- "TF" → 3 (Third Floor)
- Numeric (e.g., "4") → 4, 5, etc.
Args:
image_name (str): Image filename (e.g., "FF part 1upE.png")
Returns:
int: Floor number (e.g., 1, 2, 3, etc.)
"""
# Remove extension and split by spaces
name_no_ext = os.path.splitext(os.path.basename(image_name))[0]
parts = name_no_ext.split()
if not parts:
print(f"Warning: Could not extract floor from '{image_name}', defaulting to floor 1")
return 1
first_word = parts[0].upper()
# Map floor codes to numeric
floor_mapping = {
"FF": 1, # First Floor (lowest)
"SF": 2, # Second Floor
"TF": 3, # Third Floor
}
if first_word in floor_mapping:
floor_num = floor_mapping[first_word]
print(f"Detected floor from '{image_name}': '{first_word}' → floor {floor_num}")
return floor_num
# Check if first word is numeric
if first_word.isdigit():
floor_num = int(first_word)
print(f"Detected floor from '{image_name}': numeric '{first_word}' → floor {floor_num}")
return floor_num
# Fallback: try to extract number from first word
import re
match = re.search(r'(\d+)', first_word)
if match:
floor_num = int(match.group(1))
print(f"Detected floor from '{image_name}': extracted '{floor_num}' → floor {floor_num}")
return floor_num
# Default fallback
print(f"Warning: Could not determine floor from '{image_name}' (first word: '{first_word}'), defaulting to floor 1")
return 1
def make_graph(image_name, floor_id=None, progress_callback=None):
"""
Main function to check image existence, construct paths, run text detection,
and save graph-related outputs (plot and JSON).
Args:
image_name (str): Name of the image with the extension.
floor_id (int, optional): Floor number (e.g., 1, 2, 3).
If None, will be detected from filename or default to 1.
progress_callback (callable, optional): Called with (stage_name: str) at each pipeline stage.
"""
def _report(stage):
if progress_callback:
progress_callback(stage)
# Define base paths
base_path = os.getcwd()
input_images_dir = os.path.join(base_path, "Input_Images")
model_weights_dir = os.path.join(base_path, "Model_weights")
results_dir = os.path.join(base_path, "Results")
time_dir = os.path.join(results_dir, "Time&Meta")
os.makedirs(time_dir, exist_ok=True)
time_dir_txt = os.path.join(time_dir, "Text files")
os.makedirs(time_dir_txt, exist_ok=True)
time_dir_plots = os.path.join(time_dir, "Time Correlation Plots")
os.makedirs(time_dir_plots, exist_ok=True)
# Ensure the image exists
image_path = check_image_exists(image_name, input_images_dir)
image_name_no_ext = os.path.splitext(os.path.basename(image_path))[0]
file_size_bytes = os.path.getsize(image_path)
with Image.open(image_path) as img:
width, height = img.size
timer_file = f"{time_dir_txt}/{image_name_no_ext}_timer_info.txt"
with open(timer_file, "w") as tf:
tf.write("Timer & Metadata Information\n")
tf.write("=================\n")
tf.write(f"File Size: {file_size_bytes / 1024:.2f} KB\n")
tf.write(f"Image Dimensions: {width} x {height} pixels\n")
# Helper to log time
def log_time(stage, start_time):
elapsed_time = time.time() - start_time
with open(timer_file, "a") as tf:
tf.write(f"{stage}: {elapsed_time:.2f} seconds\n")
start_total = time.time()
cuda_available = torch.cuda.is_available()
print(f"CUDA available: {cuda_available}")
# Construct paths
plots_dir = os.path.join(results_dir, "Plots")
json_dir = os.path.join(results_dir, "Json")
json_img_dir = os.path.join(json_dir, f"{image_name_no_ext}")
os.makedirs(json_img_dir, exist_ok=True)
graph_plot_dir = os.path.join(plots_dir, "graph_plots")
os.makedirs(graph_plot_dir, exist_ok=True)
graph_img_dir = os.path.join(graph_plot_dir, f"{image_name_no_ext}")
os.makedirs(graph_img_dir, exist_ok=True)
text_detection_dir = os.path.join(plots_dir, "text_detection")
os.makedirs(text_detection_dir, exist_ok=True)
connective_plot_dir = os.path.join(plots_dir, "connective_plots")
os.makedirs(connective_plot_dir, exist_ok=True)
connect_img_dir = os.path.join(connective_plot_dir, f"{image_name_no_ext}")
os.makedirs(connect_img_dir, exist_ok=True)
test_plot_dir = os.path.join(plots_dir, "test_plots")
os.makedirs(test_plot_dir, exist_ok=True)
test_img_dir = os.path.join(test_plot_dir, f"{image_name_no_ext}")
os.makedirs(test_img_dir, exist_ok=True)
# Call get_Textboxes to perform text detection
_report("Detecting text regions")
start_step = time.time()
text_file_path = get_Textboxes(image_path, model_weights_dir, text_detection_dir)
log_time("text detection check", start_step)
_report("Interpreting text labels")
start_step = time.time()
print("\nInterpreting bboxes...")
room_bboxes, hallway_bboxes, outside_bboxes, transition_bboxes, result_file_path = interpret_bboxes(image_path, text_file_path, plots_dir)
log_time("Interpreting bboxes check", start_step)
# Initialize the graph
_report("Initializing graph nodes")
print("\nInitializing Graph")
# Detect floor from filename if not provided
if floor_id is None:
floor_id = detect_floor_from_filename(image_name)
else:
print(f"Using provided floor_id: {floor_id}")
# Convert floor_id to string for BuildingGraph (it stores as string internally)
floor_id_str = str(floor_id)
graph = BuildingGraph(default_floor=floor_id_str)
print(f"Graph initialized with default_floor: {floor_id_str} (floor number: {floor_id})")
# Calculate bounding box centers
start_step = time.time()
bbox_centers = graph.calculate_bbox_centers(room_bboxes)
hallway_bbox_centers = graph.calculate_bbox_centers(hallway_bboxes)
outside_bbox_centers = graph.calculate_bbox_centers(outside_bboxes)
transition_bbox_centers = graph.calculate_bbox_centers(transition_bboxes)
# Add nodes to the graph
for i, (x, y) in enumerate(bbox_centers):
node_id = f"room_{i + 1}"
graph.add_node(node_id, node_type="room", position=(x, y))
print(f"\nAdded {len(bbox_centers)} room nodes to the graph")
for i, (x, y) in enumerate(hallway_bbox_centers):
node_id = f"corridor_main_{i + 1}"
graph.add_node(node_id, node_type="corridor", position=(x, y))
print(f"Added {len(hallway_bbox_centers)} corridor nodes to the graph")
for i, (x, y) in enumerate(outside_bbox_centers):
node_id = f"outside_main_{i + 1}"
graph.add_node(node_id, node_type="outside", position=(x, y))
print(f"Added {len(outside_bbox_centers)} outside nodes to the graph")
transition_label_by_bbox = parse_transition_labels(result_file_path)
stairs_count = 0
elevator_count = 0
for i, ((x, y), bbox) in enumerate(zip(transition_bbox_centers, transition_bboxes), start=1):
key = tuple(bbox)
label = transition_label_by_bbox.get(key)
if label == "stairs":
stairs_count += 1
node_id = f"stairs_{stairs_count}"
elif label == "elevator":
elevator_count += 1
node_id = f"elevator_{elevator_count}"
else:
node_id = f"transition_{i}"
graph.add_node(
node_id,
node_type="transition",
position=(x, y),
floor_id=graph.default_floor,
)
print(f"Added {stairs_count} stair nodes and {elevator_count} elevator nodes to the graph\n")
log_time("graph initialization check", start_step)
json_output_path = os.path.join(json_img_dir, f"{image_name_no_ext}_ini_graph.json")
graph.save_to_json(json_output_path)
graph_plot_output_path = os.path.join(graph_img_dir, f"{image_name_no_ext}_ini_graph.png")
graph.plot_on_image(image_path, graph_plot_output_path, display_labels=True, threshold_radius = 20, highlight_regions=False)
graph_plot_output_path = os.path.join(graph_img_dir, f"{image_name_no_ext}_w_thr_graph.png")
graph.plot_on_image(image_path, graph_plot_output_path, display_labels=True, threshold_radius = 20, highlight_regions=True)
start_step = time.time()
graph.merge_nearby_nodes(threshold_room=50, threshold_door=20)
log_time("Merging nodes check", start_step)
json_output_path = os.path.join(json_img_dir, f"{image_name_no_ext}_final_graph.json")
graph.save_to_json(json_output_path)
graph_plot_output_path = os.path.join(graph_img_dir, f"{image_name_no_ext}_thr_graph.png")
graph.plot_on_image(image_path, graph_plot_output_path, display_labels=True, threshold_radius = 20, highlight_regions=False)
_report("Flood filling rooms")
print("\nFloodfilling rooms prior to door detection")
start_step = time.time()
#smart_fill_rooms(image_path, graph, results_dir, radius_threshold=50, node_radius=10)
smart_output_img_pth, smart_area_file_pth = process_fill_rooms(image_path, graph, results_dir, radius_threshold=70, node_radius=10, fill_mode="smart", point_radius=10, point_step=10, flood_threshold=30)
flood_output_img_pth, flood_area_file_pth = process_fill_rooms(image_path, graph, results_dir, radius_threshold=70, node_radius=10, fill_mode="flood", point_radius=40, point_step=10, flood_threshold=30)
log_time("Flood Filling check", start_step)
print("\nFinding Pixelwise areas")
room_pixels, outdoor_pixels, corridor_pixels, unmarked_pixels, wall_pixels, thr_img_path = pixelwise_areas(flood_output_img_pth, graph, connect_img_dir,print_tag=False)
candidates, overlay_path, room_props, area_map_path = get_room_subnode_candidates(
image_path,
graph,
results_dir=results_dir,
segmented_map_path=thr_img_path, # <<< new (strongly recommended)
spacing_px=60, wall_pad_px=10, jitter_px=4,
fill_mode="flood", point_radius=10, point_step=10, flood_threshold=30,
radius_threshold=70, save_overlay=True
)
# When adding subnodes (unchanged), also copy room_props onto MAIN room node only:
for room_id, pts in candidates.items():
floor = graph.graph.nodes[room_id].get("floor", graph.default_floor)
# attach properties to the main room node
if room_id in room_props:
graph.graph.nodes[room_id]["room_area_px"] = room_props[room_id]["area_px"]
graph.graph.nodes[room_id]["room_eq_radius"] = room_props[room_id]["eq_radius_px"]
graph.graph.nodes[room_id]["room_inradius"] = room_props[room_id]["inradius_px"]
graph.graph.nodes[room_id]["room_num_subnode_candidates"] = room_props[room_id]["num_candidates"]
graph.graph.nodes[room_id]["room_centroid_xy"] = room_props[room_id]["centroid_xy"]
for i, (x, y) in enumerate(pts, start=1):
sub_id = f"{room_id}_subnode_{i}"
graph.add_node(sub_id, node_type="room", position=(x, y), floor_id=floor)
graph.graph.nodes[sub_id]["is_subnode"] = True
graph.graph.nodes[sub_id]["parent_room_id"] = room_id
_report("Detecting doors")
print("\nDetecting doors")
start_step = time.time()
door_bbox = detect_doors(image_path,threshold=0.9, chunk_size=300, overlap=75, results_dir=plots_dir)
print("\nRefining doors")
door_bbox = refine_door_bboxes(image_path, plots_dir, door_threshold= 20, door_bboxes=door_bbox)
log_time("Detecting doors check", start_step)
_report("Classifying doors")
print("\nClassifying doors")
start_step = time.time()
exit_dbboxes, corridor2corridor_dbboxes, room2corridor_dbboxes, room2room_dbboxes, wardrobe_dbboxes = classify_doors(thr_img_path, door_bbox, connect_img_dir, print_tag=False)
log_time("Classifying doors check", start_step)
_report("Building room-door connectivity")
print("\nAdding classified door nodes to graph")
start_step = time.time()
graph.add_door_nodes(exit_dbboxes, corridor2corridor_dbboxes, room2corridor_dbboxes, room2room_dbboxes)
print("\nAdding room to door edges to graph")
graph.make_room_door_edges(image_path, (room2corridor_dbboxes+room2room_dbboxes+exit_dbboxes))
_report("Populating corridor network")
print("\nAdding corridor nodes to graph")
corridor_distance = 20
corridor_pixels = graph.add_corridor_nodes(image_path, corridor_pixels, test_img_dir, dest="corridor", distance=corridor_distance)
print(f"Added {len(corridor_pixels)} corridor nodes to the graph")
for i, (y, x) in enumerate(corridor_pixels):
node_id = f"corridor_connect_{i + 1}"
node_type = "corridor"
graph.add_node(node_id, node_type=node_type, position=(x, y))
graph.add_corridor_edges(corridor_pixels, distance=corridor_distance)
print("\nAdding outdoor nodes to graph")
outside_distance = 40
outdoor_pixels = graph.add_corridor_nodes(image_path, outdoor_pixels, test_img_dir, dest="outside", distance=outside_distance)
print(f"Added {len(outdoor_pixels)} outside nodes to the graph")
for i, (y, x) in enumerate(outdoor_pixels):
node_id = f"outside_connect_{i + 1}"
node_type = "outside"
graph.add_node(node_id, node_type=node_type, position=(x, y))
graph.add_outdoor_edges(outdoor_pixels, distance=outside_distance)
log_time("Updating graph nodes check", start_step)
_report("Funneling room paths to doors")
print("\nFunneling room families to doors (grid lattice -> shortest paths)...")
kept = graph.connect_all_families_funnel(spacing_px=60, door_selector="nearest")
print(f"Kept {kept} intra-room edges across all rooms.")
_report("Creating edges")
start_step = time.time()
graph.connect_hallways()
graph.connect_doors()
graph.connect_rooms()
graph.connect_transitions()
log_time("Edge creation check", start_step)
graph_plot_output_path = os.path.join(graph_img_dir, f"{image_name_no_ext}_pre_pruning_wothr_connect_graph.png")
graph.plot_on_image(image_path, graph_plot_output_path, display_labels=False, threshold_radius = 20, highlight_regions=False)
graph_plot_output_path = os.path.join(graph_img_dir, f"{image_name_no_ext}_pre_pruning_wthr_connect_graph.png")
graph.plot_on_image(image_path, graph_plot_output_path, display_labels=False, threshold_radius = 20, highlight_regions=True)
json_output_path = os.path.join(json_img_dir, f"{image_name_no_ext}_pre_pruning.json")
graph.save_to_json(json_output_path)
json_file_size_kb_bfr = os.path.getsize(json_output_path) / 1024
tot_graph_nodes = graph.return_graph_size()
_report("Pruning graph")
start_step = time.time()
graph.connect_all_rooms(image_path, graph_img_dir)
log_time("Graph pruning check", start_step)
mod_graph_nodes = graph.return_graph_size()
graph_plot_output_path = os.path.join(graph_img_dir, f"{image_name_no_ext}_post_pruning_wothr_connect_graph.png")
graph.plot_on_image(image_path, graph_plot_output_path, display_labels=False, threshold_radius = 20, highlight_regions=False)
graph_plot_output_path = os.path.join(graph_img_dir, f"{image_name_no_ext}_post_pruning_wthr_connect_graph.png")
graph.plot_on_image(image_path, graph_plot_output_path, display_labels=False, threshold_radius = 20, highlight_regions=True)
white_image_path = os.path.join(os.path.dirname(image_path), "white_background.png")
original_image = Image.open(image_path)
width, height = original_image.size
white_image = Image.new("RGB", (width, height), (255, 255, 255))
white_image.save(white_image_path)
graph_plot_output_path = os.path.join(graph_img_dir, f"{image_name_no_ext}_post_pruning_blank_connect_graph_1.png")
graph.plot_on_image(white_image_path, graph_plot_output_path, display_labels=False, threshold_radius=20, highlight_regions=True)
graph_plot_output_path = os.path.join(graph_img_dir, f"{image_name_no_ext}_post_pruning_blank_connect_graph_2.png")
graph.plot_on_image(white_image_path, graph_plot_output_path, display_labels=False, threshold_radius=20, highlight_regions=False)
graph_plot_output_path = os.path.join(graph_img_dir, f"{image_name_no_ext}_post_pruning_blank_connect_graph_3.png")
graph.plot_on_image(white_image_path, graph_plot_output_path, display_labels=True, threshold_radius=20, highlight_regions=False)
json_output_path = os.path.join(json_img_dir, f"{image_name_no_ext}_post_pruning.json")
graph.save_to_json(json_output_path)
json_file_size_kb_aft = os.path.getsize(json_output_path) / 1024
total_time = time.time() - start_total
with open(timer_file, "a") as tf:
tf.write(f"Total Time: {total_time:.2f} seconds\n")
tf.write(f"Total graph nodes (before pruning): {tot_graph_nodes}\n")
tf.write(f"JSON File Size (before pruning): {json_file_size_kb_bfr:.2f} KB\n")
tf.write(f"Total graph nodes (after pruning): {mod_graph_nodes}\n")
tf.write(f"JSON File Size (after pruning): {json_file_size_kb_aft:.2f} KB\n")
df_results = analyze_timer_files(time_dir_txt, time_dir_plots)
return graph, image_name_no_ext, floor_id # floor_id is int
# ============================================================================
# MULTI-FLOOR PROCESSING FUNCTIONS (READY FOR USE)
# ============================================================================
def check_results_exist(image_name, results_dir):
"""
Check if results directory exists for a given image.
Args:
image_name (str): Name of the image with extension
results_dir (str): Path to Results directory
Returns:
bool: True if results directory exists, False otherwise
"""
image_name_no_ext = os.path.splitext(os.path.basename(image_name))[0]
json_dir = os.path.join(results_dir, "Json", image_name_no_ext)
return os.path.exists(json_dir) and os.path.isdir(json_dir)
def validate_transition_mapping(transition_mapping, input_images_dir, results_dir):
"""
Validate transition mapping before processing.
Checks:
1. All images exist in Input_Images directory
2. All nodes are transition nodes (stairs/elevators)
3. Floor order constraints (no skipping floors)
4. One-to-one constraint (one transition on floor X can only connect to one transition on floor Y)
Args:
transition_mapping (dict): Transition mapping to validate
input_images_dir (str): Path to Input_Images directory
results_dir (str): Path to Results directory
Returns:
tuple: (is_valid, errors, warnings)
- is_valid (bool): True if mapping is valid
- errors (list): List of error messages
- warnings (list): List of warning messages
"""
errors = []
warnings = []
if not transition_mapping:
return True, errors, warnings
# Collect all unique images and floors from mapping
all_images = set()
all_floors = set()
node_references = {} # Track node references: (floor, image, node_id)
for (src_floor, src_image, src_node_id), targets in transition_mapping.items():
all_images.add(src_image)
all_floors.add(src_floor)
node_references[(src_floor, src_image, src_node_id)] = "source"
for tgt_floor, tgt_image, tgt_node_id in targets:
all_images.add(tgt_image)
all_floors.add(tgt_floor)
node_references[(tgt_floor, tgt_image, tgt_node_id)] = "target"
# Check 1: All images exist
print("\n" + "="*70)
print("VALIDATING TRANSITION MAPPING")
print("="*70)
print("\n[1/5] Checking image existence...")
missing_images = []
for image_name in all_images:
image_path = os.path.join(input_images_dir, image_name)
if not os.path.exists(image_path):
missing_images.append(image_name)
errors.append(f"Image '{image_name}' not found in '{input_images_dir}'")
if missing_images:
print(f" ✗ ERROR: {len(missing_images)} image(s) not found:")
for img in missing_images:
print(f" - {img}")
else:
print(f" ✓ All {len(all_images)} images exist")
# Check 2: Floor order constraints
print("\n[2/5] Checking floor order constraints...")
sorted_floors = sorted(all_floors)
# Check for direct connections between non-adjacent floors
floor_connections = {} # {floor: set of connected floors}
for (src_floor, src_image, src_node_id), targets in transition_mapping.items():
if src_floor not in floor_connections:
floor_connections[src_floor] = set()
for tgt_floor, tgt_image, tgt_node_id in targets:
floor_connections[src_floor].add(tgt_floor)
# Check if floors are adjacent (difference of 1)
floor_diff = abs(tgt_floor - src_floor)
if floor_diff > 1:
warnings.append(
f"Non-adjacent floor connection: Floor {src_floor} → Floor {tgt_floor} "
f"(difference: {floor_diff}). This may violate building semantics."
)
print(f" ⚠ WARNING: Floor {src_floor} connects directly to Floor {tgt_floor} (non-adjacent)")
if not warnings:
print(f" ✓ All floor connections are adjacent")
# Check 3: One-to-one constraint
print("\n[3/5] Checking one-to-one constraint...")
# For each floor pair (src, tgt), ensure only one transition connects them
floor_pair_transitions = {} # {(src_floor, tgt_floor): [list of transition pairs]}
for (src_floor, src_image, src_node_id), targets in transition_mapping.items():
for tgt_floor, tgt_image, tgt_node_id in targets:
pair = (src_floor, tgt_floor)
if pair not in floor_pair_transitions:
floor_pair_transitions[pair] = []
floor_pair_transitions[pair].append((src_node_id, tgt_node_id))
one_to_one_violations = []
for (src_floor, tgt_floor), transitions in floor_pair_transitions.items():
if len(transitions) > 1:
one_to_one_violations.append((src_floor, tgt_floor, transitions))
errors.append(
f"One-to-one violation: Floor {src_floor} → Floor {tgt_floor} has {len(transitions)} "
f"transition connections: {transitions}. Only one transition can connect a floor pair."
)
if one_to_one_violations:
print(f" ✗ ERROR: {len(one_to_one_violations)} one-to-one constraint violation(s):")
for src_floor, tgt_floor, transitions in one_to_one_violations:
print(f" Floor {src_floor} → Floor {tgt_floor}: {len(transitions)} connections")
for src_node, tgt_node in transitions:
print(f" - {src_node}{tgt_node}")
else:
print(f" ✓ One-to-one constraint satisfied for all floor pairs")
# Check 4: Node type validation (will be done after graphs are loaded)
print("\n[4/5] Node type validation will be performed after graphs are loaded")
# Check 5: Image name consistency
print("\n[5/5] Checking image name consistency...")
# Verify that images match their detected floors
image_floor_mismatches = []
for image_name in all_images:
detected_floor = detect_floor_from_filename(image_name)
# Find which floor this image is referenced as in the mapping
referenced_floors = set()
for (src_floor, src_image, _), targets in transition_mapping.items():
if src_image == image_name:
referenced_floors.add(src_floor)
for tgt_floor, tgt_image, _ in targets:
if tgt_image == image_name:
referenced_floors.add(tgt_floor)
# Check if detected floor matches any referenced floor
if referenced_floors and detected_floor not in referenced_floors:
image_floor_mismatches.append((image_name, detected_floor, referenced_floors))
warnings.append(
f"Image '{image_name}' detected as Floor {detected_floor}, but referenced as "
f"Floor(s) {referenced_floors} in mapping"
)
if image_floor_mismatches:
print(f" ⚠ WARNING: {len(image_floor_mismatches)} image-floor mismatch(es):")
for img, detected, referenced in image_floor_mismatches:
print(f" {img}: detected={detected}, referenced={referenced}")
else:
print(f" ✓ All image-floor references are consistent")
is_valid = len(errors) == 0
print(f"\n{'='*70}")
if is_valid:
print("✓ VALIDATION PASSED")
if warnings:
print(f" ({len(warnings)} warning(s) - see above)")
else:
print("✗ VALIDATION FAILED")
print(f" {len(errors)} error(s) found - see above")
print(f"{'='*70}\n")
return is_valid, errors, warnings
def ensure_floor_graph_exists(image_name, input_images_dir, results_dir):
"""
Ensure that a floor graph exists for the given image.
If image exists but graph doesn't, process it first.
Args:
image_name (str): Name of the image with extension
input_images_dir (str): Path to Input_Images directory
results_dir (str): Path to Results directory
Returns:
tuple: (floor_graph, floor_num, image_path) or (None, None, None) if image doesn't exist
"""
# Check if image exists
image_path = os.path.join(input_images_dir, image_name)
if not os.path.exists(image_path):
print(f" ✗ ERROR: Image '{image_name}' not found in '{input_images_dir}'")
return None, None, None
# Detect floor
floor_num = detect_floor_from_filename(image_name)
# Check if results exist
if check_results_exist(image_name, results_dir):
print(f" ✓ Results exist for '{image_name}' (Floor {floor_num})")
# Load existing graph (we'll process it fresh anyway to ensure consistency)
# For now, we'll always reprocess to ensure graphs are up-to-date
print(f" → Processing '{image_name}' to create/update graph...")
floor_graph, image_name_no_ext, detected_floor = make_graph(image_name, floor_id=floor_num)
return floor_graph, floor_num, image_path
else:
print(f" → No results found for '{image_name}' (Floor {floor_num})")
print(f" → Processing '{image_name}' to create graph...")
floor_graph, image_name_no_ext, detected_floor = make_graph(image_name, floor_id=floor_num)
return floor_graph, floor_num, image_path
def normalize_image_coordinates(image_path, position):
'''
Normalize coordinates to account for different image scales.
Returns normalized (x, y) in [0, 1] range.
Args:
image_path (str): Path to the image
position (tuple): (x, y) pixel coordinates
Returns:
tuple: Normalized (x_norm, y_norm) coordinates
'''
with Image.open(image_path) as img:
width, height = img.size
x_norm = position[0] / width if width > 0 else 0.0
y_norm = position[1] / height if height > 0 else 0.0
return (x_norm, y_norm)
def align_transitions_spatially(transition1_pos, transition2_pos, image1_path, image2_path,
tolerance=0.02):
'''
Check if two transitions are vertically aligned (same x,y position regardless of scale).
Uses normalized coordinates for scale-invariant comparison.
Args:
transition1_pos (tuple): (x, y) position of transition in image1
transition2_pos (tuple): (x, y) position of transition in image2
image1_path (str): Path to first image
image2_path (str): Path to second image
tolerance (float): Maximum allowed difference in normalized coordinates (default: 0.02 = 2%)
Returns:
bool: True if transitions are vertically aligned
'''
norm1 = normalize_image_coordinates(image1_path, transition1_pos)
norm2 = normalize_image_coordinates(image2_path, transition2_pos)
dx = abs(norm1[0] - norm2[0])
dy = abs(norm1[1] - norm2[1])
is_aligned = (dx <= tolerance) and (dy <= tolerance)
if is_aligned:
print(f" ✓ Transitions aligned: ({norm1[0]:.4f}, {norm1[1]:.4f}) vs ({norm2[0]:.4f}, {norm2[1]:.4f})")
print(f" Difference: dx={dx:.4f}, dy={dy:.4f} (tolerance={tolerance})")
else:
print(f" ✗ Transitions NOT aligned: ({norm1[0]:.4f}, {norm1[1]:.4f}) vs ({norm2[0]:.4f}, {norm2[1]:.4f})")
print(f" Difference: dx={dx:.4f}, dy={dy:.4f} (tolerance={tolerance})")
return is_aligned
def merge_floor_graphs(floor_graphs, floor_image_paths):
'''
Merge multiple floor graphs into a single graph, preserving floor information.
Node IDs are prefixed with floor number to avoid conflicts.
Args:
floor_graphs (dict): {floor_num: BuildingGraph} - Graphs for each floor (floor_num is int)
floor_image_paths (dict): {floor_num: image_path} - Image paths for each floor
Returns:
BuildingGraph: Merged graph with all floors
'''
print("\n" + "="*70)
print("MERGING MULTI-FLOOR GRAPHS")
print("="*70)
merged_graph = BuildingGraph(default_floor="MULTI_FLOOR")
total_nodes = 0
total_edges = 0
# Sort floors by floor number for consistent processing
sorted_floors = sorted(floor_graphs.keys())
print(f"Processing floors in order: {sorted_floors}")
# Add all nodes from all floors with floor prefix
for floor_num in sorted_floors:
floor_graph = floor_graphs[floor_num]
floor_id_str = str(floor_num) # Convert to string for node prefix
print(f"\nProcessing floor {floor_num}:")
floor_nodes = 0
floor_edges = len(floor_graph.graph.edges())
for node_id, node_data in floor_graph.graph.nodes(data=True):
# Create unique node ID with floor prefix (e.g., "1_room_1", "2_stairs_1")
prefixed_node_id = f"{floor_num}_{node_id}"
# Preserve original floor information (store as string in graph)
merged_graph.add_node(
prefixed_node_id,
node_type=node_data['type'],
position=node_data['position'],
floor_id=floor_id_str # Store as string in graph
)
# Copy all other attributes
for key, value in node_data.items():
if key not in ['type', 'position', 'floor']:
merged_graph.graph.nodes[prefixed_node_id][key] = value
floor_nodes += 1
# Add edges within this floor (with prefixed node IDs)
for u, v, edge_data in floor_graph.graph.edges(data=True):
prefixed_u = f"{floor_num}_{u}"
prefixed_v = f"{floor_num}_{v}"
if merged_graph.graph.has_node(prefixed_u) and merged_graph.graph.has_node(prefixed_v):
merged_graph.add_edge(prefixed_u, prefixed_v, weight=edge_data.get('weight'))
# Copy edge attributes
for key, value in edge_data.items():
if key not in ['weight', 'distance']:
merged_graph.graph[prefixed_u][prefixed_v][key] = value
total_nodes += floor_nodes
total_edges += floor_edges
print(f" Added {floor_nodes} nodes and {floor_edges} edges from floor {floor_num}")
print(f"\nMerged graph summary:")
print(f" Total nodes: {total_nodes}")
print(f" Total edges: {total_edges}")
print(f" Floors: {sorted_floors}")
return merged_graph
def connect_transitions_across_floors(merged_graph, floor_graphs, floor_image_paths,
transition_mapping, spatial_tolerance=0.02):
'''
Connect transition nodes across floors based on manual mapping and spatial alignment.
Args:
merged_graph (BuildingGraph): Merged multi-floor graph
floor_graphs (dict): {floor_num: BuildingGraph} - Original floor graphs (floor_num is int)
floor_image_paths (dict): {floor_num: image_path} - Image paths for each floor
transition_mapping (dict): Manual mapping of transitions
Format: {
(source_floor_num, source_image, source_node_id): [
(target_floor_num, target_image, target_node_id),
...
]
}
Example: {
(1, "FF part 1upE.png", "stairs_1"): [
(2, "SF part 1upE.png", "stairs_1"),
(3, "TF part 1upE.png", "stairs_1")
]
}
Note: floor_num is integer (1, 2, 3, etc.)
spatial_tolerance (float): Tolerance for spatial alignment check (default: 0.02)
Returns:
int: Number of inter-floor connections created
'''
print("\n" + "="*70)
print("CONNECTING TRANSITIONS ACROSS FLOORS")
print("="*70)
connections_created = 0
connections_failed = 0
# Track connections per floor pair to enforce one-to-one constraint
floor_pair_connections = {} # {(src_floor, tgt_floor): (src_node_id, tgt_node_id)}
for (src_floor_num, src_image, src_node_id), targets in transition_mapping.items():
print(f"\nProcessing transition mapping:")
print(f" Source: Floor {src_floor_num} / {src_image} / {src_node_id}")
# Get source node from original floor graph
if src_floor_num not in floor_graphs:
print(f" ✗ ERROR: Floor {src_floor_num} not found in floor_graphs")
print(f" Available floors: {list(floor_graphs.keys())}")
connections_failed += len(targets)
continue
src_graph = floor_graphs[src_floor_num]
if src_node_id not in src_graph.graph.nodes:
print(f" ✗ ERROR: Node '{src_node_id}' not found in floor {src_floor_num} graph")
# List available transition nodes
transition_nodes = [n for n in src_graph.graph.nodes()
if src_graph.graph.nodes[n].get('type') == 'transition']
if transition_nodes:
print(f" Available transition nodes: {transition_nodes[:10]}")
else:
print(f" No transition nodes found on this floor")
connections_failed += len(targets)
continue
src_node_data = src_graph.graph.nodes[src_node_id]
src_node_type = src_node_data.get('type')
# Validate node type - must be transition
if src_node_type != 'transition':
print(f" ✗ ERROR: Source node '{src_node_id}' is not a transition node (type: '{src_node_type}')")
print(f" Only transition nodes (stairs/elevators) can be connected across floors")
connections_failed += len(targets)
continue
src_pos = src_node_data.get('position')
src_image_path = floor_image_paths.get(src_floor_num)
if src_pos is None:
print(f" ✗ ERROR: Source node '{src_node_id}' has no position")
connections_failed += len(targets)
continue
if src_image_path is None:
print(f" ✗ ERROR: No image path found for floor {src_floor_num}")
print(f" Available floors: {list(floor_image_paths.keys())}")
connections_failed += len(targets)
continue
src_prefixed_id = f"{src_floor_num}_{src_node_id}"
if src_prefixed_id not in merged_graph.graph.nodes:
print(f" ✗ ERROR: Prefixed source node '{src_prefixed_id}' not found in merged graph")
connections_failed += len(targets)
continue
# Connect to each target
for tgt_floor_num, tgt_image, tgt_node_id in targets:
print(f" Target: Floor {tgt_floor_num} / {tgt_image} / {tgt_node_id}")
# Check floor order constraint
floor_diff = abs(tgt_floor_num - src_floor_num)
if floor_diff == 0:
print(f" ✗ ERROR: Cannot connect node to same floor ({src_floor_num})")
connections_failed += 1
continue
if floor_diff > 1:
print(f" ⚠ WARNING: Non-adjacent floor connection (Floor {src_floor_num} → Floor {tgt_floor_num})")
print(f" This violates building semantics - floors should connect sequentially")
# Continue anyway but warn
# Get target node from original floor graph
if tgt_floor_num not in floor_graphs:
print(f" ✗ ERROR: Floor {tgt_floor_num} not found in floor_graphs")
print(f" Available floors: {list(floor_graphs.keys())}")
connections_failed += 1
continue
tgt_graph = floor_graphs[tgt_floor_num]
if tgt_node_id not in tgt_graph.graph.nodes:
print(f" ✗ ERROR: Node '{tgt_node_id}' not found in floor {tgt_floor_num} graph")
# List available transition nodes
transition_nodes = [n for n in tgt_graph.graph.nodes()
if tgt_graph.graph.nodes[n].get('type') == 'transition']
if transition_nodes:
print(f" Available transition nodes: {transition_nodes[:10]}")
else:
print(f" No transition nodes found on this floor")
connections_failed += 1
continue
tgt_node_data = tgt_graph.graph.nodes[tgt_node_id]
tgt_node_type = tgt_node_data.get('type')
# Validate node type - must be transition
if tgt_node_type != 'transition':
print(f" ✗ ERROR: Target node '{tgt_node_id}' is not a transition node (type: '{tgt_node_type}')")
print(f" Only transition nodes (stairs/elevators) can be connected across floors")
connections_failed += 1
continue
tgt_pos = tgt_node_data.get('position')
tgt_image_path = floor_image_paths.get(tgt_floor_num)
if tgt_pos is None:
print(f" ✗ ERROR: Target node '{tgt_node_id}' has no position")
connections_failed += 1
continue
if tgt_image_path is None:
print(f" ✗ ERROR: No image path found for floor {tgt_floor_num}")
print(f" Available floors: {list(floor_image_paths.keys())}")
connections_failed += 1
continue
tgt_prefixed_id = f"{tgt_floor_num}_{tgt_node_id}"
if tgt_prefixed_id not in merged_graph.graph.nodes:
print(f" ✗ ERROR: Prefixed target node '{tgt_prefixed_id}' not found in merged graph")
connections_failed += 1
continue
# Enforce one-to-one constraint: check if this floor pair already has a connection
floor_pair = (src_floor_num, tgt_floor_num)
if floor_pair in floor_pair_connections:
existing_src, existing_tgt = floor_pair_connections[floor_pair]
print(f" ✗ ERROR: One-to-one constraint violation!")
print(f" Floor {src_floor_num} → Floor {tgt_floor_num} already connected via:")
print(f" {existing_src}{existing_tgt}")
print(f" Cannot add another connection: {src_node_id}{tgt_node_id}")
print(f" Each floor pair can only have ONE transition connection")
connections_failed += 1
continue
# Record this connection
floor_pair_connections[floor_pair] = (src_node_id, tgt_node_id)
# Check spatial alignment
print(f" Checking spatial alignment...")
is_aligned = align_transitions_spatially(
src_pos, tgt_pos, src_image_path, tgt_image_path,
tolerance=spatial_tolerance
)
if not is_aligned:
print(f" ⚠ WARNING: Transitions are not spatially aligned!")
print(f" Proceeding anyway based on manual mapping...")
# Create inter-floor edge
if merged_graph.graph.has_edge(src_prefixed_id, tgt_prefixed_id):
print(f" ⚠ Edge already exists between {src_prefixed_id} and {tgt_prefixed_id}")
else:
# Calculate distance (use normalized coordinates for consistency)
src_norm = normalize_image_coordinates(src_image_path, src_pos)
tgt_norm = normalize_image_coordinates(tgt_image_path, tgt_pos)
# Use a fixed weight for inter-floor transitions (e.g., 1000 units)
# This represents vertical movement between floors
inter_floor_weight = 1000.0
merged_graph.add_edge(
src_prefixed_id,
tgt_prefixed_id,
weight=inter_floor_weight
)
# Mark as inter-floor edge
merged_graph.graph[src_prefixed_id][tgt_prefixed_id]['edge_type'] = 'inter_floor'
merged_graph.graph[src_prefixed_id][tgt_prefixed_id]['spatial_aligned'] = is_aligned
print(f" ✓ Created inter-floor connection: {src_prefixed_id}{tgt_prefixed_id}")
print(f" Weight: {inter_floor_weight} (inter-floor transition)")
connections_created += 1
print(f"\nInter-floor connection summary:")
print(f" Connections created: {connections_created}")
print(f" Connections failed: {connections_failed}")
return connections_created
def process_multi_floor(image_names, transition_mapping=None, spatial_tolerance=0.02):
'''
Process multiple floorplan images and create a unified multi-floor graph.
Args:
image_names (list): List of image filenames (e.g., ["FF part 1upE.png", "SF part 1upE.png"])
transition_mapping (dict, optional): Manual mapping of transitions across floors
Format: {
(source_floor_num, source_image, source_node_id): [
(target_floor_num, target_image, target_node_id),
...
]
}
Example: {
(1, "FF part 1upE.png", "stairs_1"): [
(2, "SF part 1upE.png", "stairs_1"),
(3, "TF part 1upE.png", "stairs_1")
]
}
Note: floor_num is integer (1, 2, 3, etc.)
spatial_tolerance (float): Tolerance for spatial alignment check (default: 0.02)
Returns:
tuple: (merged_graph, floor_graphs, floor_image_paths)
- merged_graph: Unified BuildingGraph with all floors
- floor_graphs: Dict {floor_num: BuildingGraph} of individual floor graphs (floor_num is int)
- floor_image_paths: Dict {floor_num: image_path} mapping floor number to image path
'''
print("\n" + "="*70)
print("MULTI-FLOOR PROCESSING")
print("="*70)
base_path = os.getcwd()
input_images_dir = os.path.join(base_path, "Input_Images")
results_dir = os.path.join(base_path, "Results")
# Step 1: Validate transition mapping if provided
if transition_mapping:
is_valid, errors, warnings = validate_transition_mapping(
transition_mapping, input_images_dir, results_dir
)
if not is_valid:
print("\n✗ TRANSITION MAPPING VALIDATION FAILED")
print("Cannot proceed with invalid mapping. Please fix the errors above.")
raise ValueError(f"Transition mapping validation failed with {len(errors)} error(s)")
if warnings:
print(f"\n⚠ {len(warnings)} warning(s) found - proceeding anyway")
# Collect all unique images from mapping
all_mapping_images = set()
for (src_floor, src_image, _), targets in transition_mapping.items():
all_mapping_images.add(src_image)
for tgt_floor, tgt_image, _ in targets:
all_mapping_images.add(tgt_image)
# Ensure all images from mapping are in image_names
missing_from_list = all_mapping_images - set(image_names)
if missing_from_list:
print(f"\n⚠ Adding {len(missing_from_list)} image(s) from mapping to processing list:")
for img in missing_from_list:
print(f" - {img}")
image_names = list(set(image_names) | missing_from_list)
print(f"\nProcessing {len(image_names)} floorplan images...")
start_multi_total = time.time()
floor_graphs = {}
floor_image_paths = {}
# Step 2: Process each floor (ensure graphs exist)
for image_name in image_names:
print(f"\n{'='*70}")
print(f"Processing image: {image_name}")
print(f"{'='*70}")
start_floor = time.time()
# Ensure graph exists (will process if needed)
floor_graph, floor_num, image_path = ensure_floor_graph_exists(
image_name, input_images_dir, results_dir
)
if floor_graph is None:
print(f" ✗ ERROR: Failed to process '{image_name}' - skipping")
continue
# Store results (use floor_num as key)
floor_graphs[floor_num] = floor_graph
floor_image_paths[floor_num] = image_path
floor_time = time.time() - start_floor
print(f"\nCompleted processing floor {floor_num} in {floor_time:.2f} seconds")
print(f" Nodes: {floor_graph.return_graph_size()}")
print(f" Edges: {len(floor_graph.graph.edges())}")
# Merge all floor graphs
print(f"\n{'='*70}")
print("MERGING FLOOR GRAPHS")
print(f"{'='*70}")
start_merge = time.time()
merged_graph = merge_floor_graphs(floor_graphs, floor_image_paths)
merge_time = time.time() - start_merge
print(f"Merging completed in {merge_time:.2f} seconds")
# Connect transitions across floors
if transition_mapping:
print(f"\n{'='*70}")
print("CONNECTING TRANSITIONS ACROSS FLOORS")
print(f"{'='*70}")
start_connect = time.time()
connections = connect_transitions_across_floors(
merged_graph, floor_graphs, floor_image_paths,
transition_mapping, spatial_tolerance=spatial_tolerance
)
connect_time = time.time() - start_connect
print(f"Inter-floor connection completed in {connect_time:.2f} seconds")
print(f" Created {connections} inter-floor connections")
else:
print("\n⚠ No transition mapping provided - skipping inter-floor connections")
print(" Provide transition_mapping to connect floors via transitions")
# Save merged graph
json_dir = os.path.join(results_dir, "Json")
merged_json_dir = os.path.join(json_dir, "MULTI_FLOOR")
os.makedirs(merged_json_dir, exist_ok=True)
merged_json_path = os.path.join(merged_json_dir, "merged_multi_floor_graph.json")
merged_graph.save_to_json(merged_json_path)
json_size_kb = os.path.getsize(merged_json_path) / 1024
print(f"\nMerged graph saved to: {merged_json_path}")
print(f" JSON size: {json_size_kb:.2f} KB")
print(f" Total nodes: {merged_graph.return_graph_size()}")
print(f" Total edges: {len(merged_graph.graph.edges())}")
total_time = time.time() - start_multi_total
print(f"\n{'='*70}")
print(f"MULTI-FLOOR PROCESSING COMPLETE")
print(f"{'='*70}")
print(f"Total processing time: {total_time:.2f} seconds")
print(f" Per-floor processing: {total_time - merge_time - (connect_time if transition_mapping else 0):.2f} seconds")
print(f" Merging: {merge_time:.2f} seconds")
if transition_mapping:
print(f" Inter-floor connections: {connect_time:.2f} seconds")
return merged_graph, floor_graphs, floor_image_paths
# Example usage (commented out - ready for use when multi-floor images are available):
"""
# Example: Process 3 floors and connect transitions
image_names = [
"FF part 1upE.png", # Floor 1
"SF part 1upE.png", # Floor 2
"TF part 1upE.png" # Floor 3
]
# Manual transition mapping (floor numbers are integers: 1, 2, 3, etc.)
# Format: (source_floor_num, source_image, source_node_id): [(target_floor_num, target_image, target_node_id), ...]
transition_mapping = {
# stairs_1 connects Floor 1 → Floor 2 → Floor 3
(1, "FF part 1upE.png", "stairs_1"): [
(2, "SF part 1upE.png", "stairs_1"),
(3, "TF part 1upE.png", "stairs_1")
],
# stairs_2 connects Floor 1 → Floor 2 → Floor 3
(1, "FF part 1upE.png", "stairs_2"): [
(2, "SF part 1upE.png", "stairs_2"),
(3, "TF part 1upE.png", "stairs_2")
],
# elevator_1 connects all floors
(1, "FF part 1upE.png", "elevator_1"): [
(2, "SF part 1upE.png", "elevator_1"),
(3, "TF part 1upE.png", "elevator_1")
]
}
# Process multi-floor
merged_graph, floor_graphs, floor_image_paths = process_multi_floor(
image_names,
transition_mapping=transition_mapping,
spatial_tolerance=0.02
)
"""
if __name__ == "__main__":
# Set up argument parser
parser = argparse.ArgumentParser(
description="Run text detection on a given image and generate graph outputs.",
epilog="Example usage: python main.py image_name.png"
)
parser.add_argument(
"image_name",
type=str,
nargs="?",
help="Name of the image (with extension) in the Input_Images folder."
)
# Parse arguments
args = parser.parse_args()
# Check if image_name is provided
if not args.image_name:
parser.error("The following argument is required: image_name (name of the image with extension)")
try:
# Run the main function
make_graph(args.image_name)
except FileNotFoundError as e:
# Print the error and exit
print(e)
sys.exit(1)