File size: 10,077 Bytes
1e3df84 ef4c917 1e3df84 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 | """
Processing pipeline wrapper for the Tesseract++ system
Wraps Main.py functionality for web API usage
"""
import os
import sys
import json
import tempfile
import shutil
from pathlib import Path
from typing import Dict, Any, Optional
import threading
from typing import Callable
# Headless matplotlib backend before Main (and its plotting modules) load.
os.environ.setdefault("MPLBACKEND", "Agg")
import matplotlib
matplotlib.use("Agg")
# Progress tracking (thread-safe via GIL for simple dict writes)
_progress: Dict[str, str] = {}
def set_progress(key: str, stage: str):
_progress[key] = stage
def get_progress(key: str) -> str:
return _progress.get(key, "")
def clear_progress(key: str):
_progress.pop(key, None)
# Add required paths
base_path = Path(__file__).parent.parent.parent.parent # Back to Tesseract++ root
sys.path.insert(0, str(base_path))
sys.path.insert(0, str(base_path / "Models" / "Text_Models"))
sys.path.insert(0, str(base_path / "Models" / "Interpreter"))
sys.path.insert(0, str(base_path / "Models" / "Door_Models"))
sys.path.insert(0, str(base_path / "utils"))
# Import main processing functions
import Main
class TimeoutException(Exception):
"""Custom exception for processing timeout"""
pass
class ProcessingPipeline:
"""
Wrapper for the Tesseract++ processing pipeline
"""
def __init__(self):
self.base_path = Path(__file__).parent.parent.parent.parent
self.input_images_dir = self.base_path / "Input_Images"
self.results_dir = self.base_path / "Results"
self.temp_dir = self.base_path / "temp_processing"
# Create temp directory if not exists
self.temp_dir.mkdir(exist_ok=True)
# Verify model weights exist
self._verify_models()
def _verify_models(self):
"""Verify all required model weights are present"""
weights_dir = self.base_path / "Model_weights"
required_weights = [
"craft_mlt_25k.pth",
"None-VGG-BiLSTM-CTC.pth",
"door_mdl_32.pth"
]
for weight_file in required_weights:
weight_path = weights_dir / weight_file
if not weight_path.exists():
raise FileNotFoundError(f"Required model weight not found: {weight_path}")
def get_cached_result(self, image_name: str) -> Optional[Dict[str, Any]]:
"""
Check for and return pre-computed results for an image.
Args:
image_name: Image filename (e.g. "FF part 1upE.png")
Returns:
Processing result dict if cached, None otherwise.
"""
image_stem = Path(image_name).stem
json_dir = self.results_dir / "Json" / image_stem
post_pruning_json = json_dir / f"{image_stem}_post_pruning.json"
pre_pruning_json = json_dir / f"{image_stem}_pre_pruning.json"
if not post_pruning_json.exists():
return None
with open(post_pruning_json, 'r') as f:
graph_data = json.load(f)
stats = self._calculate_statistics(graph_data)
# Load pre-pruning graph if available
pre_pruning_graph = None
if pre_pruning_json.exists():
with open(pre_pruning_json, 'r') as f:
pre_pruning_graph = json.load(f)
pre_nodes = len(pre_pruning_graph.get("nodes", []))
post_nodes = len(graph_data.get("nodes", []))
stats["pruning_reduction"] = round(
(1 - post_nodes / pre_nodes) * 100, 2
) if pre_nodes > 0 else 0
return {
"graph_json": graph_data,
"pre_pruning_graph_json": pre_pruning_graph,
"stats": stats,
"image_name": image_name
}
def has_cached_result(self, image_name: str) -> bool:
"""Check if a cached result exists for an image."""
image_stem = Path(image_name).stem
post_pruning_json = self.results_dir / "Json" / image_stem / f"{image_stem}_post_pruning.json"
return post_pruning_json.exists()
def process_image(self, image_path: str, image_name: str, timeout: int = 180, progress_key: str = None) -> Dict[str, Any]:
"""
Process a floorplan image through the Tesseract++ pipeline
Args:
image_path: Path to the image file
image_name: Original image filename
timeout: Processing timeout in seconds
Returns:
Dictionary containing processing results
"""
# Create a unique temp folder for this processing session
import uuid
session_id = str(uuid.uuid4())
session_dir = self.temp_dir / session_id
session_dir.mkdir(exist_ok=True)
# Copy image to Input_Images temporarily if not already there
input_image_path = self.input_images_dir / image_name
image_was_copied = False
# Track exception from thread
thread_exception = [None]
try:
if not input_image_path.exists():
shutil.copy2(image_path, input_image_path)
image_was_copied = True
# Redirect outputs to session directory
original_cwd = os.getcwd()
os.chdir(self.base_path)
# Progress callback
def on_progress(stage: str):
if progress_key:
set_progress(progress_key, stage)
# Run the main processing pipeline with timeout
def run_processing():
try:
Main.make_graph(image_name, progress_callback=on_progress)
except Exception as e:
thread_exception[0] = e
# Use threading for timeout control
thread = threading.Thread(target=run_processing)
thread.daemon = True
thread.start()
thread.join(timeout)
if thread.is_alive():
raise TimeoutException(f"Processing exceeded {timeout} seconds")
if thread_exception[0] is not None:
raise thread_exception[0]
# Extract results
image_name_no_ext = Path(image_name).stem
# Find the generated JSON files
json_dir = self.results_dir / "Json" / image_name_no_ext
post_pruning_json = json_dir / f"{image_name_no_ext}_post_pruning.json"
pre_pruning_json = json_dir / f"{image_name_no_ext}_pre_pruning.json"
# Read the post-pruning graph
if not post_pruning_json.exists():
raise FileNotFoundError(f"Processing completed but output not found: {post_pruning_json}")
with open(post_pruning_json, 'r') as f:
graph_data = json.load(f)
# Calculate statistics
stats = self._calculate_statistics(graph_data)
# Load pre-pruning graph if available
pre_pruning_graph = None
if pre_pruning_json.exists():
with open(pre_pruning_json, 'r') as f:
pre_pruning_graph = json.load(f)
pre_nodes = len(pre_pruning_graph.get("nodes", []))
post_nodes = len(graph_data.get("nodes", []))
stats["pruning_reduction"] = round((1 - post_nodes / pre_nodes) * 100, 2) if pre_nodes > 0 else 0
result = {
"graph_json": graph_data,
"pre_pruning_graph_json": pre_pruning_graph,
"stats": stats,
"session_id": session_id,
"image_name": image_name
}
# Clean up Results for uploaded (non-example) images
if image_was_copied:
self._cleanup_results(image_name_no_ext)
return result
finally:
# Cleanup
os.chdir(original_cwd)
if progress_key:
clear_progress(progress_key)
# Remove copied image if it was temporary
if image_was_copied and input_image_path.exists():
input_image_path.unlink()
# Clean up session directory
if session_dir.exists():
shutil.rmtree(session_dir, ignore_errors=True)
def _cleanup_results(self, image_stem: str):
"""Clean up Results subdirectories for non-example (uploaded) images."""
results_subdirs = [
"Json", "Plots/connective_plots", "Plots/door_detect",
"Plots/flood_fill", "Plots/graph_plots", "Plots/interpreter_detect",
"Plots/room_subnodes", "Plots/smart_fill", "Plots/text_detection",
"Plots/test_plots", "Time&Meta/Text files"
]
for subdir in results_subdirs:
result_path = self.results_dir / subdir / image_stem
if result_path.exists() and result_path.is_dir():
shutil.rmtree(result_path, ignore_errors=True)
# Also check for timer info text files
timer_file = self.results_dir / "Time&Meta" / "Text files" / f"{image_stem}_timer_info.txt"
if timer_file.exists():
timer_file.unlink(missing_ok=True)
def _calculate_statistics(self, graph_data: Dict[str, Any]) -> Dict[str, Any]:
"""Calculate graph statistics"""
nodes = graph_data.get("nodes", [])
edges = graph_data.get("edges", [])
# Count nodes by type
node_types = {}
for node in nodes:
node_type = node.get("type", "unknown")
node_types[node_type] = node_types.get(node_type, 0) + 1
return {
"total_nodes": len(nodes),
"total_edges": len(edges),
"node_types": node_types
}
def get_example_images(self) -> list:
"""Get list of available example images"""
images = []
for img_path in sorted(self.input_images_dir.glob("*.png"))[:4]:
images.append({
"name": img_path.name,
"size_kb": img_path.stat().st_size / 1024
})
return images
|