import io import os import tempfile from collections import Counter import multiprocessing import numpy as np import pandas as pd import streamlit as st from PIL import Image import job_store # With multiprocessing "spawn", every worker process re-imports __main__ (this # file). We must not execute any Streamlit calls inside a worker or they will # produce ScriptRunContext warnings and can cause the worker to hang. # All helper functions below are safe to define in both contexts; only the UI # execution block at the bottom is gated on _IS_MAIN. _IS_MAIN = multiprocessing.current_process().name == "MainProcess" if _IS_MAIN: # Diagnostic marker: fires on every single script execution, including # Streamlit reruns and fresh sessions from new WebSocket connections. # Comparing session_id across consecutive prints tells us whether the # analysis loop is being killed by an in-session rerun (same session_id) # or the browser/proxy dropping and re-opening the WebSocket connection # (different session_id each time) — see restart-loop investigation. import time as _diag_time from streamlit.runtime.scriptrunner import get_script_run_ctx as _get_script_run_ctx _diag_ctx = _get_script_run_ctx() _diag_session_id = _diag_ctx.session_id if _diag_ctx else "NO_CTX" _diag_run_count = st.session_state.get("_diag_run_count", 0) + 1 st.session_state["_diag_run_count"] = _diag_run_count print( f"[app][{_diag_time.strftime('%H:%M:%S')}] SCRIPT START #{_diag_run_count} " f"pid={os.getpid()} session_id={_diag_session_id}", flush=True, ) # ── Helpers ─────────────────────────────────────────────────────────────────── def _bgr_pages_from_bytes(pdf_bytes: bytes, dpi: int = 200) -> list[np.ndarray]: """Render each PDF page to a BGR numpy array using PyMuPDF (no poppler needed).""" import fitz # PyMuPDF doc = fitz.open(stream=pdf_bytes, filetype="pdf") mat = fitz.Matrix(dpi / 72, dpi / 72) pages = [] for page in doc: pix = page.get_pixmap(matrix=mat, colorspace=fitz.csRGB) arr = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, 3) pages.append(arr[:, :, ::-1].copy()) # RGB → BGR return pages def _pages_to_pdf_bytes(bgr_images: list[np.ndarray]) -> bytes: """Convert BGR numpy arrays back to a PDF via Pillow (no cv2 needed).""" pils = [Image.fromarray(img[:, :, ::-1]) for img in bgr_images] # BGR → RGB buf = io.BytesIO() pils[0].save(buf, format="PDF", save_all=True, append_images=pils[1:]) return buf.getvalue() def _pages_matching_text(pdf_bytes: bytes, needle: str, crop_dpi: int = 150) -> list[dict]: """Search each page's bottom-right quadrant for `needle` (case-insensitive). Architectural title blocks live there, so restricting the search to that corner skips boilerplate notes and callouts scattered across the rest of the sheet that happen to mention the same words. Matching works at word granularity (via PyMuPDF's "words" extraction) rather than a plain substring check, since title-block sheet names often wrap across two lines (e.g. "1st Story Floor" / "Plan Code Study") -- concatenating word-by-word with single spaces (instead of the newline PyMuPDF would join lines with) lets "floor plan" match across that wrap, and keeping each word's own rect lets the match be highlighted precisely even when it spans lines. Returns one dict per matching page -- {"page": 1-indexed page number, "snippet": surrounding text, "crop": RGB uint8 array of the matched region with a box drawn around the match} -- so a human can see what was actually matched before trusting it.""" import fitz # PyMuPDF from PIL import ImageDraw CONTEXT = 50 # characters of snippet context on each side of the match CROP_PAD_PT = 40 # points of padding around the match, for the preview crop BOX_COLOR = (220, 30, 30) doc = fitz.open(stream=pdf_bytes, filetype="pdf") needle_norm = " ".join(needle.lower().split()) scale = crop_dpi / 72 matches = [] for i, page in enumerate(doc): r = page.rect quadrant = fitz.Rect(r.width / 2, r.height / 2, r.width, r.height) words = page.get_text("words", clip=quadrant) if not words: continue # Concatenate words (PyMuPDF's own reading-order) with single spaces, # tracking each word's character span so a match can be mapped back # to the word rect(s) it came from. concat = "" spans = [] # (start_char, end_char, word_rect) for w in words: if concat: concat += " " start = len(concat) concat += w[4] spans.append((start, len(concat), fitz.Rect(w[0], w[1], w[2], w[3]))) pos = concat.lower().find(needle_norm) if pos == -1: continue match_end = pos + len(needle_norm) involved = [rect for s, e, rect in spans if e > pos and s < match_end] if not involved: continue match_rect = involved[0] for rect in involved[1:]: match_rect |= rect # Grow the box 50% (25% per side) around its own center so # descenders/ascenders and characters that PyMuPDF's word rects clip # a little tight on are still fully inside the drawn box. _grow_x = match_rect.width * 0.25 _grow_y = match_rect.height * 0.25 box_rect = fitz.Rect( match_rect.x0 - _grow_x, match_rect.y0 - _grow_y, match_rect.x1 + _grow_x, match_rect.y1 + _grow_y, ) snip_start = max(0, pos - CONTEXT) snip_end = min(len(concat), match_end + CONTEXT) snippet = concat[snip_start:snip_end] if snip_start > 0: snippet = "…" + snippet if snip_end < len(concat): snippet = snippet + "…" crop_rect = fitz.Rect( max(0, box_rect.x0 - CROP_PAD_PT), max(0, box_rect.y0 - CROP_PAD_PT), min(r.width, box_rect.x1 + CROP_PAD_PT), min(r.height, box_rect.y1 + CROP_PAD_PT), ) pix = page.get_pixmap(matrix=fitz.Matrix(scale, scale), clip=crop_rect) crop_img = Image.frombytes("RGB", (pix.width, pix.height), pix.samples) ImageDraw.Draw(crop_img).rectangle( ( (box_rect.x0 - crop_rect.x0) * scale, (box_rect.y0 - crop_rect.y0) * scale, (box_rect.x1 - crop_rect.x0) * scale, (box_rect.y1 - crop_rect.y0) * scale, ), outline=BOX_COLOR, width=3, ) matches.append({ "page": i + 1, "snippet": snippet, "crop": np.array(crop_img), }) doc.close() return matches def _extract_pdf_pages(pdf_bytes: bytes, page_indices: list[int]) -> bytes: """Return a new PDF (as bytes) containing only the given 0-indexed pages, in the order given.""" import fitz # PyMuPDF doc = fitz.open(stream=pdf_bytes, filetype="pdf") doc.select(page_indices) out = doc.tobytes() doc.close() return out def _annotate_pages(results: list[dict], valid) -> list[np.ndarray]: import classical_cv_detector as ccv return [ccv.annotate(r["bgr"], r["dets"], valid=valid) for r in results] def _find_text_matches( pdf_bytes: bytes, needle: str, dpi: int = 200 ) -> list[list[tuple[int, int, int, int]]]: """ Search every page's embedded text layer for `needle` (case-insensitive) and return, per page, the pixel-space bbox of each occurrence -- scaled by dpi/72 to line up with the page images from _bgr_pages_from_bytes, which uses the same default dpi. Scanned/rasterized pages with no text layer will simply yield zero matches. """ import fitz # PyMuPDF scale = dpi / 72 doc = fitz.open(stream=pdf_bytes, filetype="pdf") matches = [ [ (int(r.x0 * scale), int(r.y0 * scale), int(r.x1 * scale), int(r.y1 * scale)) for r in page.search_for(needle) ] for page in doc ] doc.close() return matches def _mark_text_matches( pages: list[np.ndarray], matches_per_page: list[list[tuple[int, int, int, int]]] ) -> list[np.ndarray]: """Draw a box around every text match, on a copy of each page.""" import cv2 CLR_TEXT_MATCH = (255, 255, 0) # cyan (BGR) -- distinct from annotate()'s # green/red/blue/magenta detection colours out = [] for page, boxes in zip(pages, matches_per_page): img = page.copy() for x1, y1, x2, y2 in boxes: cv2.rectangle(img, (x1, y1), (x2, y2), CLR_TEXT_MATCH, 3) out.append(img) return out def _build_clean_bytes(results: list[dict], annotated_pages: list[np.ndarray]) -> bytes: """ One PDF page per input page: all detection crops arranged in a left-to-right, top-to-bottom grid. Crops are taken from the annotated source so the green/red/blue colour coding is preserved. """ import cv2 CELL_MARGIN = 12 # px gap around each crop LABEL_H = 22 # px reserved below each crop for the code label COLS = 6 # columns per row clean_pages = [] for r, ann in zip(results, annotated_pages): ph, pw = ann.shape[:2] crops: list[tuple[np.ndarray, str]] = [] for d in r["dets"]: x, y, bw, bh = d["bbox"] pad = 15 x1 = max(0, x - pad) y1 = max(0, y - pad) x2 = min(pw, x + bw + pad) y2 = min(ph, y + bh + pad) if x2 > x1 and y2 > y1: crops.append((ann[y1:y2, x1:x2].copy(), d["code"])) if not crops: clean_pages.append(np.full((200, 800, 3), 255, dtype=np.uint8)) continue max_cw = max(c.shape[1] for c, _ in crops) max_ch = max(c.shape[0] for c, _ in crops) cols = min(COLS, len(crops)) cell_w = max_cw + 2 * CELL_MARGIN cell_h = max_ch + LABEL_H + 2 * CELL_MARGIN rows = (len(crops) + cols - 1) // cols canvas_w = cols * cell_w + CELL_MARGIN canvas_h = rows * cell_h + CELL_MARGIN canvas = np.full((canvas_h, canvas_w, 3), 255, dtype=np.uint8) for i, (crop, code) in enumerate(crops): row_i, col_i = divmod(i, cols) cell_x = CELL_MARGIN + col_i * cell_w cell_y = CELL_MARGIN + row_i * cell_h ch, cw = crop.shape[:2] x_off = (max_cw - cw) // 2 y_off = (max_ch - ch) // 2 canvas[cell_y + y_off : cell_y + y_off + ch, cell_x + x_off : cell_x + x_off + cw] = crop cv2.rectangle( canvas, (cell_x + x_off - 1, cell_y + y_off - 1), (cell_x + x_off + cw, cell_y + y_off + ch), (200, 200, 200), 1, ) (tw, _), _ = cv2.getTextSize(code, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1) tx = cell_x + (max_cw - tw) // 2 ty = cell_y + max_ch + CELL_MARGIN + 14 cv2.putText(canvas, code, (tx, ty), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (60, 60, 60), 1, cv2.LINE_AA) clean_pages.append(canvas) return _pages_to_pdf_bytes(clean_pages) def _build_csv_bytes(results: list[dict], legend_df=None) -> bytes: rows = [] for r in results: codes = [d["code"] for d in r["dets"] if d["code"] != "?"] for code, count in Counter(codes).items(): rows.append({"Page": r["page"], "Code": code, "Count": count}) if not rows: return b"Code,Total\n" df = pd.DataFrame(rows) pivot = df.pivot_table( index="Code", columns="Page", values="Count", fill_value=0, aggfunc="sum", ) pivot.columns = [f"Page {c}" for c in pivot.columns] pivot["Total"] = pivot.sum(axis=1) pivot = pivot.reset_index() if legend_df is not None and "TYPE MARK" in legend_df.columns: legend_cols = [c for c in legend_df.columns if c != "TYPE MARK"] legend_trim = legend_df[["TYPE MARK"] + legend_cols].rename( columns={"TYPE MARK": "Code"} ) pivot = legend_trim.merge(pivot, on="Code", how="right") return pivot.to_csv(index=False).encode() def _summary_pivot(results: list[dict], valid) -> pd.DataFrame | None: rows = [] for r in results: codes = [d["code"] for d in r["dets"] if d["code"] != "?"] unknown = sum(1 for d in r["dets"] if d["code"] == "?") for code, cnt in Counter(codes).items(): rows.append({"Page": r["page"], "Code": code, "Count": cnt}) if unknown: rows.append({"Page": r["page"], "Code": "?", "Count": unknown}) if not rows: return None df = pd.DataFrame(rows) pivot = df.pivot_table( index="Code", columns="Page", values="Count", fill_value=0, aggfunc="sum", ) pivot.columns = [f"Page {c}" for c in pivot.columns] pivot["Total"] = pivot.sum(axis=1) if valid is not None: pivot.index = pd.Index( [ f"{c} [valid]" if c in valid else (c if c == "?" else f"{c} [flagged]") for c in pivot.index ], name="Code", ) total_row = pivot.sum(axis=0).rename("Total").to_frame().T total_row.index.name = "Code" return pd.concat([pivot, total_row]) def _get_validated_results(results: list[dict]) -> list[dict]: """ Return a copy of results reflecting the current state of the Validate tab widgets. Unchecked detections are dropped; edited codes are substituted. Falls back to the original detection when a widget key hasn't been created yet (i.e. the Validate tab has never been opened). """ validated = [] for r in results: validated_dets = [] for j, d in enumerate(r["dets"]): key_check = f"det_check_{r['page']}_{j}" key_code = f"det_code_{r['page']}_{j}" included = st.session_state.get(key_check, True) code = st.session_state.get(key_code, d["code"]) if included: validated_dets.append({**d, "code": code}) validated.append({**r, "dets": validated_dets}) return validated # ── Streamlit UI — only runs in the main process ────────────────────────────── if _IS_MAIN: st.set_page_config( page_title="Blueprint Window Shape Detector", layout="wide", initial_sidebar_state="collapsed", ) # ── Session state ───────────────────────────────────────────────────────── _STATE_DEFAULTS = { "running": False, "pending": None, # {pdf, legend_bytes, legend_name} written before rerun "results": None, "valid": None, "legend_df": None, "ann_bytes": None, "csv_bytes": None, "clean_bytes": None, "ann_pages": None, # list[np.ndarray] — kept for Validate tab crops "last_run_summary": None, # {total_min, pg_done, n, failed} — survives past running=False "floor_plan_source_id": None, # blueprint_file.file_id the search below was run against "floor_plan_matches": None, # [{"page": 1-idx, "snippet": str, "crop": ndarray}, ...] from the last search "floor_plan_pdf_bytes": None, # filtered PDF bytes built from the applied selection, or None "floor_plan_applied_pages": None, # 1-indexed page numbers baked into floor_plan_pdf_bytes "mode": "Run Detection Analysis", # or "Search Blueprint Text" "search_source_id": None, # blueprint_file.file_id the search results below were computed from "search_results": {}, # {search text: [count on page 1, count on page 2, ...]} "search_marked_pdf": None, # bytes — all searched terms highlighted, built on demand } for _key, _val in _STATE_DEFAULTS.items(): if _key not in st.session_state: st.session_state[_key] = _val if not st.session_state.running and not st.session_state.pending: _orphan_key = job_store.any_active_key() if _orphan_key is not None: # A job is still running server-side in job_store from a # previous session this browser tab has no memory of (a brand # new session gets fresh, empty st.session_state, but job_store # is a plain module-level dict that survives session loss -- see # job_store.py). Resume showing progress for it instead of # silently landing on the idle upload screen while the job # keeps running unseen. st.session_state.running = True st.session_state.pending = {"_orphan_job_key": _orphan_key} # ── Tabs ────────────────────────────────────────────────────────────────── _ready = bool(st.session_state.results) tab_input, tab_search, tab_progress, tab_output, tab_validate = st.tabs([ "Input", "Search", "Progress" if st.session_state.running else "Progress (idle)", "Output" if _ready else "Output (locked)", "Validate" if _ready else "Validate (locked)", ]) # ── PROGRESS TAB ────────────────────────────────────────────────────────── # Widgets live inside the tab (not above it) so the main page stays clean # while an analysis runs -- previously this whole block rendered above # every tab, on screen no matter which tab was open. st.empty() # placeholders update in place regardless of which tab is currently # selected in the browser, so the polling loop further down still works # unchanged; the user just needs this tab open to see the updates. _progress_bar = None _status_text = None _page_status_container = None with tab_progress: if st.session_state.running: st.info("Analysis in progress. Please wait — this may take several minutes per page.") _progress_bar = st.progress(0) _status_text = st.empty() _page_status_container = st.empty() elif st.session_state.last_run_summary: # Show the completed/failed run's final numbers instead of just # going back to a bare idle message -- previously this # information (total runtime, per-page breakdown) was computed # and then immediately discarded on the st.rerun() that follows # completion, since running flips to False right before it. _summary = st.session_state.last_run_summary if _summary.get("failed"): st.error(f"Last analysis failed after {_summary['total_min']} minutes.") else: st.success(f"Last analysis completed in {_summary['total_min']} minutes.") _lines = ["**Page Breakdown**"] for _pi in range(_summary["n"]): if _pi in _summary["pg_done"]: _lines.append( f"    Page {_pi + 1} finished processing " f"in {_summary['pg_done'][_pi]} minutes" ) st.markdown(" \n".join(_lines)) st.caption("Start a new analysis from the **Input** tab.") else: st.caption("No analysis is currently running. Start one from the **Input** tab.") # ══ INPUT TAB ════════════════════════════════════════════════════════════ with tab_input: st.title("Blueprint Diamond Detector") st.caption( "Detects diamond-shaped unit symbols (e.g. A1, B2) in construction blueprints " "using classical computer vision and OCR." ) st.divider() blueprint_file = st.file_uploader( "Blueprint PDF *(required)*", type=["pdf"], help="The construction blueprint PDF containing diamond unit symbols.", disabled=bool(st.session_state.running), ) if blueprint_file: import fitz as _fitz _doc = _fitz.open(stream=blueprint_file.read(), filetype="pdf") _n_pages = len(_doc) _doc.close() blueprint_file.seek(0) _parallel_batches = (_n_pages + 1) // 2 _est_lo = _parallel_batches * 5 _est_hi = _parallel_batches * 8 _time_note = ( "approximately 5–8 minutes" if _n_pages == 1 else f"approximately {_est_lo}–{_est_hi} minutes" ) st.caption( f"**{_n_pages} page(s) detected.** " f"Pages are analysed 2 at a time in parallel (using both vCPUs), " f"so the estimated runtime is {_time_note}." ) st.radio( "What would you like to do?", ["Run Detection Analysis", "Search Blueprint Text"], key="mode", horizontal=True, disabled=bool(st.session_state.running), ) if st.session_state.mode == "Search Blueprint Text": st.info( "Detection analysis is disabled while Search is selected. " "Switch to the **Search** tab above to search this blueprint's text." ) else: if st.session_state.floor_plan_source_id != blueprint_file.file_id: # A different file than the one the current search (if any) # was run against -- drop the stale results rather than # silently applying them to unrelated pages. st.session_state.floor_plan_source_id = None st.session_state.floor_plan_matches = None st.session_state.floor_plan_pdf_bytes = None st.session_state.floor_plan_applied_pages = None if st.button( "Search for 'Floor Plan' Pages", disabled=bool(st.session_state.running), ): st.session_state.floor_plan_source_id = blueprint_file.file_id st.session_state.floor_plan_matches = _pages_matching_text( blueprint_file.getvalue(), "Floor Plan" ) st.session_state.floor_plan_pdf_bytes = None st.session_state.floor_plan_applied_pages = None for _k in list(st.session_state.keys()): if _k.startswith("floor_plan_check_"): del st.session_state[_k] st.rerun() if st.session_state.floor_plan_source_id == blueprint_file.file_id: _fp_matches = st.session_state.floor_plan_matches if not _fp_matches: st.warning("No pages containing the text 'Floor Plan' were found.") else: _fp_included = sum( st.session_state.get(f"floor_plan_check_{m['page']}", True) for m in _fp_matches ) with st.expander( f"Floor Plan Matches — {len(_fp_matches)} of {_n_pages} page(s), " f"{_fp_included} included", expanded=False, ): st.caption( "Each page below matched 'Floor Plan' in its title-block " "corner. Uncheck any that aren't real floor plan pages, " "then press **Apply Selection** to commit." ) with st.form(key="floor_plan_form", enter_to_submit=False): for _m in _fp_matches: _key_check = f"floor_plan_check_{_m['page']}" _fp_c1, _fp_c2 = st.columns([1, 4]) with _fp_c1: st.checkbox( f"Page {_m['page']}", value=st.session_state.get(_key_check, True), key=_key_check, ) with _fp_c2: st.image(_m["crop"], caption=_m["snippet"], width=280) _fp_submitted = st.form_submit_button( "Apply Selection", use_container_width=True ) if _fp_submitted: _fp_selected = [ m["page"] for m in _fp_matches if st.session_state.get(f"floor_plan_check_{m['page']}", True) ] if _fp_selected: st.session_state.floor_plan_pdf_bytes = _extract_pdf_pages( blueprint_file.getvalue(), [p - 1 for p in _fp_selected], ) st.session_state.floor_plan_applied_pages = _fp_selected else: st.session_state.floor_plan_pdf_bytes = None st.session_state.floor_plan_applied_pages = None st.warning("No pages selected — all pages will be used instead.") st.rerun() if st.session_state.floor_plan_applied_pages: _fp_applied = st.session_state.floor_plan_applied_pages st.success( f"Using {len(_fp_applied)} of {_n_pages} page(s): " f"{', '.join(str(n) for n in _fp_applied)}." ) if st.button("Use All Pages Instead", disabled=bool(st.session_state.running)): st.session_state.floor_plan_pdf_bytes = None st.session_state.floor_plan_applied_pages = None st.rerun() legend_file = st.file_uploader( "Legend Image *(optional)*", type=["png", "jpg", "jpeg", "bmp", "tiff", "webp"], help="A screenshot or scan of the window schedule / legend. " "Used to validate detected codes as confirmed or flagged.", disabled=bool(st.session_state.running), ) exemplar_file = st.file_uploader( "Example Target Shape *(optional)*", type=["png", "jpg", "jpeg", "bmp", "tiff", "webp"], help="A tight crop of a single instance of the target symbol, including " "its interior code (e.g. A1). When provided, detection uses this " "shape instead of the default diamond detector.", disabled=bool(st.session_state.running), ) if exemplar_file: import cv2 as _cv2 import classical_cv_detector as _ccv _exemplar_bytes = exemplar_file.read() exemplar_file.seek(0) _exemplar_bgr = _cv2.imdecode( np.frombuffer(_exemplar_bytes, np.uint8), _cv2.IMREAD_COLOR ) _exemplar_data = _ccv.extract_shape_template(_exemplar_bgr) if _exemplar_bgr is not None else None if _exemplar_data is None: st.warning( "No clear shape outline could be found in this image. " "Try a tighter crop with the symbol's outline clearly visible." ) else: _prev_col1, _prev_col2 = st.columns(2) with _prev_col1: st.caption(f"Detected shape type: **{_exemplar_data['shape_type']}**") st.image(_exemplar_bgr[:, :, ::-1], caption="Example", width=120) with _prev_col2: st.caption("Extracted mask") st.image(_exemplar_data["mask"], caption="Mask", width=120) if st.button( "Run Analysis", type="primary", disabled=( bool(st.session_state.running) or blueprint_file is None or st.session_state.mode == "Search Blueprint Text" ), ): _use_floor_plan_filter = ( st.session_state.floor_plan_source_id == blueprint_file.file_id and st.session_state.floor_plan_pdf_bytes is not None ) st.session_state.pending = { "pdf": ( st.session_state.floor_plan_pdf_bytes if _use_floor_plan_filter else blueprint_file.read() ), "legend_bytes": legend_file.read() if legend_file else None, "legend_name": legend_file.name if legend_file else None, "exemplar_bytes": exemplar_file.read() if exemplar_file else None, } # Clear previous results and all validation widget state for _k in ("results", "valid", "legend_df", "ann_bytes", "csv_bytes", "clean_bytes", "ann_pages"): st.session_state[_k] = None for _k in list(st.session_state.keys()): if _k.startswith("det_check_") or _k.startswith("det_code_"): del st.session_state[_k] st.session_state.running = True st.rerun() # ══ SEARCH TAB ═══════════════════════════════════════════════════════════ with tab_search: st.subheader("Search Blueprint Text") if blueprint_file is None: st.warning("Upload a blueprint PDF on the **Input** tab first.") else: if st.session_state.search_source_id != blueprint_file.file_id: # A different file than the one the search results below were # computed from -- drop stale results rather than showing # counts for pages that no longer correspond to this file. st.session_state.search_source_id = blueprint_file.file_id st.session_state.search_results = {} st.session_state.search_marked_pdf = None st.caption( "Search the blueprint's embedded text layer for a phrase and see how " "many times it appears on each page. Only text in the PDF's text layer " "is searched -- scanned/rasterized pages won't have matches. Enter one " "phrase at a time; each search is added to the table below." ) with st.form(key="search_form", clear_on_submit=True): _search_col, _button_col = st.columns([4, 1]) with _search_col: _search_needle = st.text_input( "Text to search for", placeholder="e.g. Floor Plan", label_visibility="collapsed", ) with _button_col: _search_submitted = st.form_submit_button( "Search", use_container_width=True ) if _search_submitted and _search_needle.strip(): _needle = _search_needle.strip() _matches = _find_text_matches(blueprint_file.getvalue(), _needle) st.session_state.search_results[_needle] = [len(m) for m in _matches] st.session_state.search_marked_pdf = None # stale — rebuild on demand if st.session_state.search_results: import fitz as _fitz_search _search_doc = _fitz_search.open( stream=blueprint_file.getvalue(), filetype="pdf" ) _search_n_pages = len(_search_doc) _search_doc.close() _rows = [] for _term, _counts in st.session_state.search_results.items(): _row = {"Search Text": _term} for _pi in range(_search_n_pages): _row[f"Page {_pi + 1}"] = _counts[_pi] _row["Total"] = sum(_counts) _rows.append(_row) _search_df = pd.DataFrame(_rows).set_index("Search Text") st.dataframe(_search_df, use_container_width=True) if st.button("Generate Highlighted PDF (all searched terms)"): _marked_pages = _bgr_pages_from_bytes(blueprint_file.getvalue()) for _term in st.session_state.search_results: _term_matches = _find_text_matches(blueprint_file.getvalue(), _term) _marked_pages = _mark_text_matches(_marked_pages, _term_matches) st.session_state.search_marked_pdf = _pages_to_pdf_bytes(_marked_pages) if st.session_state.search_marked_pdf: st.download_button( "Download Highlighted PDF", data=st.session_state.search_marked_pdf, file_name="search_highlighted.pdf", mime="application/pdf", ) else: st.caption("No searches yet — enter text above and press **Search**.") # ══ OUTPUT TAB ═══════════════════════════════════════════════════════════ with tab_output: if not st.session_state.results: st.warning("Results are not yet available. Upload a blueprint and run the analysis on the **Input** tab first.") else: results = st.session_state.results valid = st.session_state.valid # Use validated results (reflects Validate tab edits) for summary + CSV _v_results = _get_validated_results(results) st.subheader("Detection Summary") st.caption("Updates live as you make changes in the Validate tab.") pivot = _summary_pivot(_v_results, valid) if pivot is not None: st.dataframe(pivot, use_container_width=True) if valid is not None: st.caption( "[valid] = code found in legend " "[flagged] = code not in legend " "? = shape detected but OCR failed" ) else: st.write("No detections found.") st.divider() st.subheader("Downloads") col1, col2, col3 = st.columns(3) with col1: st.download_button( label="Download Annotated PDF", data=st.session_state.ann_bytes, file_name="annotated_blueprint.pdf", mime="application/pdf", use_container_width=True, ) st.caption( "Original pages with bounding boxes. " "Green = confirmed, Red = flagged, Blue = shape-only. " "Use Regenerate PDFs in the Validate tab to reflect edits." ) with col2: st.download_button( label="Download CSV Report", data=_build_csv_bytes(_v_results, st.session_state.legend_df), file_name="detection_report.csv", mime="text/csv", use_container_width=True, ) st.caption( "Pivot table of code counts by page. Reflects Validate tab edits live." ) with col3: st.download_button( label="Download Detections-Only PDF", data=st.session_state.clean_bytes, file_name="detections_only.pdf", mime="application/pdf", use_container_width=True, ) st.caption( "Grid of detection crops per page. " "Use Regenerate PDFs in the Validate tab to reflect edits." ) # ══ VALIDATE TAB ═════════════════════════════════════════════════════════ with tab_validate: if not st.session_state.results: st.warning("Results are not yet available. Upload a blueprint and run the analysis on the **Input** tab first.") else: results = st.session_state.results valid = st.session_state.valid ann_pages = st.session_state.ann_pages # annotated BGR arrays st.subheader("Validate Detections") st.caption( "Each detection is shown with its annotated crop. " "Uncheck false positives or edit codes, then press **Apply Changes** " "for that page to commit. The summary table and CSV update immediately. " "Use **Regenerate PDFs** at the bottom to rebuild the PDF exports." ) _VALIDATE_COLS = 6 # Crop pad: enough to show the text label drawn above each bbox. # annotate() draws text at y-6 with ~15px height, so we need >21px # above the box top. 25px gives comfortable margin on all sides. _CROP_PAD = 25 for _ri, _r in enumerate(results): _page_num = _r["page"] _dets = _r["dets"] _ann_page = ann_pages[_ri] if ann_pages else None _ph, _pw = (_ann_page.shape[:2] if _ann_page is not None else _r["bgr"].shape[:2]) _src_page = _ann_page if _ann_page is not None else _r["bgr"] _n_included = sum( st.session_state.get(f"det_check_{_page_num}_{j}", True) for j in range(len(_dets)) ) with st.expander( f"Page {_page_num} — {len(_dets)} detection(s), " f"{_n_included} included", expanded=False, ): if not _dets: st.write("No detections on this page.") continue # Sort: green (confirmed) first, red (flagged) second, # blue (shape-only / "?") last. Original index is preserved # alongside each detection so widget keys remain stable. def _sort_priority(code): if code == "?": return 2 if valid is None or code in valid: return 0 return 1 _dets_sorted = sorted( enumerate(_dets), key=lambda x: _sort_priority(x[1]["code"]), ) # enter_to_submit=False prevents pressing Enter in a code # field from triggering form submission (Streamlit >= 1.39). with st.form(key=f"validate_form_{_page_num}", enter_to_submit=False): # Fixed-height scrollable container (Streamlit >= 1.30) so # opening an expander doesn't push the rest of the page down. # The caption and submit button sit below it, always visible. with st.container(height=1000): for _row_start in range(0, len(_dets_sorted), _VALIDATE_COLS): _row_dets = _dets_sorted[_row_start : _row_start + _VALIDATE_COLS] _cols = st.columns(_VALIDATE_COLS) # always full width for _ci, (_col, (_det_idx, _d)) in enumerate(zip(_cols, _row_dets)): _x, _y, _bw, _bh = _d["bbox"] _x1 = max(0, _x - _CROP_PAD) _y1 = max(0, _y - _CROP_PAD) _x2 = min(_pw, _x + _bw + _CROP_PAD) _y2 = min(_ph, _y + _bh + _CROP_PAD) _crop_rgb = _src_page[_y1:_y2, _x1:_x2, ::-1] _key_check = f"det_check_{_page_num}_{_det_idx}" _key_code = f"det_code_{_page_num}_{_det_idx}" with _col: st.image(_crop_rgb, width="stretch") st.checkbox( "Include", value=st.session_state.get(_key_check, True), key=_key_check, ) st.text_input( "Code", value=st.session_state.get(_key_code, _d["code"]), key=_key_code, label_visibility="collapsed", placeholder="e.g. A1", ) st.caption( "Uncheck any false positives and correct any misread codes above, " "then press this button to save your changes. The summary table " "and CSV in the Output tab will update immediately." ) st.form_submit_button( f"Apply Changes — Page {_page_num}", use_container_width=True, ) st.divider() st.caption( "Once you have applied changes across all pages, press this button to rebuild " "the annotated and detections-only PDFs so the downloads in the Output tab " "reflect your edits." ) if st.button("Regenerate PDFs", type="primary"): _v = _get_validated_results(st.session_state.results) _new_ann = _annotate_pages(_v, valid) st.session_state.ann_pages = _new_ann st.session_state.ann_bytes = _pages_to_pdf_bytes(_new_ann) st.session_state.clean_bytes = _build_clean_bytes(_v, _new_ann) st.success("PDFs regenerated — download updated files from the Output tab.") # ══ ANALYSIS (runs after the UI re-renders with all widgets disabled) ═════ if st.session_state.running and st.session_state.pending: import time as _time def _ts() -> str: # Wall-clock stamp on our own log lines -- Streamlit/Tornado's own # log lines already carry a timestamp, so this lets the two be # correlated directly when diagnosing a restart's exact timing. return _time.strftime("%H:%M:%S") _pending = st.session_state.pending _is_orphan_attach = "_orphan_job_key" in _pending if _is_orphan_attach: _job_key = _pending["_orphan_job_key"] else: _job_key = job_store.job_key( _pending["pdf"], _pending["legend_bytes"], _pending.get("exemplar_bytes") ) if not _is_orphan_attach and job_store.claim(_job_key): # First time seeing this exact job (same PDF+legend+exemplar # bytes): do the one-time setup (legend parse, PDF render, spawn # the Pool) and stash everything in job_store, keyed by a hash # of the input bytes rather than by Streamlit session -- see # job_store.py for why session-scoped storage isn't enough # (a session can be destroyed outright, not just reconnected). # claim() (not a plain get()-is-None check) guards this so two # reruns racing on the same job can't both create a Pool. # # Deliberately NOT using `with pool:` -- that would close the # pool the moment THIS script execution ends, but if a rerun or # a brand new session picks this job back up mid-poll, we want # it to find the pool still alive and keep polling it rather # than closing it out from under the still-running workers. _valid = None if _pending["legend_bytes"]: # Dispatch by the exemplar's classified shape_type, mirroring # classical_cv_detector.run()'s own dispatch (see its comment # at the equivalent point): hexagon schedules use a 'WINDOW # LETTER' column (single letter, or letter+0-2 digits -- e.g. # Colorado Grand Oaks' W1..W11), diamond/default schedules use # 'TYPE MARK' (letter+digit, e.g. A1). This dispatch was # already fixed in the CLI script but never ported here -- # app.py always called the diamond-only parser regardless of # exemplar shape, which raises on a hexagon-format legend # (no 'TYPE MARK' column) and was being silently swallowed # by the except below, leaving _valid=None. With no legend # to check against, annotate() then marks every OCR'd code # "confirmed" (green) -- including random page text that # happens to OCR into something letter+digit-shaped -- with # no visible error at all. import cv2 as _cv2_legend import classical_cv_detector as _ccv_legend _is_hexagon_legend = False _exemplar_bytes_for_legend = _pending.get("exemplar_bytes") if _exemplar_bytes_for_legend: _ex_bgr = _cv2_legend.imdecode( np.frombuffer(_exemplar_bytes_for_legend, np.uint8), _cv2_legend.IMREAD_COLOR ) _ex_data = _ccv_legend.extract_shape_template(_ex_bgr) if _ex_bgr is not None else None _is_hexagon_legend = _ex_data is not None and _ex_data["shape_type"] == "hexagon" _suffix = "." + _pending["legend_name"].rsplit(".", 1)[-1] with tempfile.NamedTemporaryFile(suffix=_suffix, delete=False) as _tmp: _tmp.write(_pending["legend_bytes"]) _tmp_path = _tmp.name try: if _is_hexagon_legend: from legend_parser import parse_hexagon_legend_image, valid_codes_hexagon _legend_df = parse_hexagon_legend_image(_tmp_path) _valid = valid_codes_hexagon(_legend_df) else: from legend_parser import parse_legend_image, valid_codes _legend_df = parse_legend_image(_tmp_path) _valid = valid_codes(_legend_df) st.session_state.legend_df = _legend_df except Exception: _valid = None finally: try: os.unlink(_tmp_path) except OSError: pass _status_text.text("Rendering PDF pages…") _pages = _bgr_pages_from_bytes(_pending["pdf"]) _n = len(_pages) import classical_cv_detector as ccv _ctx = multiprocessing.get_context("spawn") # Cross-process progress channel: stdout print()s from inside a # worker are only visible via HF's log-stream endpoint, which has # proven unreliable for this (a live capture across two full # 8-page runs captured zero of the tile-level progress lines # added for the pages-6/7 stall, despite ~90+ expected). A # multiprocessing.Manager().dict() proxy was tried next and ALSO # showed nothing -- and because that failure mode was silently # swallowed too, there was no way to tell whether writes were # failing or simply never happening. Plain files on the # container's local disk are about as hard to silently break as # cross-process communication gets: the main process and every # worker subprocess share the same filesystem unconditionally, # no proxy/socket lifecycle involved. One file per page index. _progress_dir = f"/tmp/tile_progress_{_job_key}" os.makedirs(_progress_dir, exist_ok=True) _exemplar_bytes = _pending.get("exemplar_bytes") _worker_args = [ (p.tobytes(), p.shape, p.dtype.str, ccv.TILE_OVERLAP, ccv.OCR_UPSCALE, _exemplar_bytes, _progress_dir, _pi) for _pi, p in enumerate(_pages) ] _pg_start: dict[int, float] = {} # page index -> start timestamp _pg_done: dict[int, int] = {} # page index -> elapsed minutes once finished _t0 = _time.time() for _pi in range(min(2, _n)): _pg_start[_pi] = _t0 print(f"[app][{_ts()}] creating Pool(processes=2) for {_n} page(s)", flush=True) _t_pool = _time.time() _pool = _ctx.Pool(processes=2, initializer=ccv.lower_worker_priority) print(f"[app][{_ts()}] pool created in {_time.time()-_t_pool:.1f}s, submitting tasks", flush=True) _async_results = [ _pool.apply_async(ccv.detect_page_worker, (arg,)) for arg in _worker_args ] job_store.put(_job_key, { "pool": _pool, "progress_dir": _progress_dir, "async_results": _async_results, "pages": _pages, "n": _n, "all_dets": [None] * _n, "pg_done": _pg_done, "pg_start": _pg_start, "t_pool": _t_pool, "valid": _valid, "legend_df": st.session_state.legend_df, "done_count": 0, }) else: # Either the job already fully exists, another thread just # claimed it and is still building it (spawning a Pool + parsing # the legend can take a few seconds), or this is a fresh session # auto-adopting an orphaned job -- wait for it to appear rather # than racing ahead with job_store.get() returning None. _wait_start = _time.time() while job_store.get(_job_key) is None and _time.time() - _wait_start < 30: _time.sleep(0.2) print(f"[app][{_ts()}] rerun/reconnect detected -- reattaching to in-flight " f"analysis (job={_job_key[:8]})", flush=True) _state = job_store.get(_job_key) if _state is None: # The job vanished (finished and was popped by another session) # in the narrow window between us discovering it and getting # here -- nothing left to attach to, so drop back to idle rather # than crash on a missing state dict. st.session_state.running = False st.session_state.pending = None st.rerun() st.session_state.legend_df = _state["legend_df"] _pool = _state["pool"] _async_results = _state["async_results"] _pages = _state["pages"] _n = _state["n"] _all_dets = _state["all_dets"] _pg_done = _state["pg_done"] _pg_start = _state["pg_start"] _t_pool = _state["t_pool"] _valid = _state["valid"] # .get() with a fallback: a job already in flight when this deploy # landed was created by prior code and has no "progress_dir" key -- # reattaching to it must not KeyError. _progress_dir = _state.get("progress_dir") def _render_page_status(): # Computed fresh on every call (rather than reading a static # cached string) so an in-progress page's elapsed time ticks up # in real time as this function is re-invoked on every poll # cycle below, not just when a page transitions to done. lines = ["**Processing Pages**"] _now = _time.time() for _pi in range(_n): if _pi in _pg_done: lines.append( f"    Page {_pi + 1} finished processing " f"in {_pg_done[_pi]} minutes" ) elif _pi in _pg_start: _live_min = round((_now - _pg_start[_pi]) / 60) lines.append( f"    Page {_pi + 1} has been processing " f"for {_live_min} minutes" ) _detail = None if _progress_dir: try: with open(os.path.join(_progress_dir, f"page_{_pi}.txt")) as _f: _detail = _f.read().strip() except Exception: # File may not exist yet (worker hasn't written # its first checkpoint) -- never let a # diagnostic read break the actual progress UI. _detail = None if _detail: lines.append(f"        → {_detail}") _page_status_container.markdown(" \n".join(lines)) _progress_bar.progress(_state["done_count"] / _n) _render_page_status() # Poll for results instead of blocking on pool.imap(): a fully blocking # wait sends zero WebSocket traffic to the browser for however long a # page takes (several minutes), and the HF Spaces reverse proxy treats # that as a dead connection and reconnects, which reruns this script # from the top (occasionally landing on a brand new session if the # old one was destroyed) -- which is why the pool/async-results are # persisted in job_store above rather than kept as local variables # that a rerun would silently discard. _POLL_INTERVAL = 2.0 # seconds between readiness checks _HEARTBEAT_INTERVAL = 15.0 # seconds between keepalive UI pushes _done_count = _state["done_count"] _last_heartbeat = _time.time() try: while _done_count < _n: for _i, _ar in enumerate(_async_results): if _all_dets[_i] is not None or not _ar.ready(): continue _dets = _ar.get() _all_dets[_i] = _dets _done_count += 1 _state["done_count"] = _done_count _elapsed_min = round((_time.time() - _pg_start.get(_i, _t_pool)) / 60) print(f"[app][{_ts()}] received page {_i + 1}/{_n} result in {_elapsed_min}min", flush=True) _progress_bar.progress(_done_count / _n) _pg_done[_i] = _elapsed_min _next = _i + 2 if _next < _n and _next not in _pg_start: _pg_start[_next] = _time.time() _render_page_status() _last_heartbeat = _time.time() if _done_count >= _n: break # Re-render on every poll cycle (not just on page-completion # transitions or the coarser heartbeat below) so in-progress # pages' elapsed-minutes text updates in real time. _render_page_status() _now = _time.time() if _now - _last_heartbeat >= _HEARTBEAT_INTERVAL: _status_text.text( f"Still processing… ({_done_count}/{_n} pages complete, " f"{round((_now - _t_pool) / 60)} min elapsed)" ) _last_heartbeat = _now _time.sleep(_POLL_INTERVAL) except Exception: # Don't leak the worker Pool on an unexpected error -- an orphaned # Pool would keep burning CPU indefinitely, which is exactly the # kind of contention this whole persistence scheme is trying to # avoid. Also clear the stored state so a retry starts clean # instead of reattaching to a dead job. print(f"[app][{_ts()}] analysis failed, terminating pool", flush=True) _pool.terminate() _pool.join() if _progress_dir: import shutil as _shutil _shutil.rmtree(_progress_dir, ignore_errors=True) job_store.pop(_job_key) st.session_state.last_run_summary = { "total_min": round((_time.time() - _t_pool) / 60), "pg_done": dict(_pg_done), "n": _n, "failed": True, } st.session_state.running = False st.session_state.pending = None raise _pool.close() _pool.join() if _progress_dir: import shutil as _shutil _shutil.rmtree(_progress_dir, ignore_errors=True) job_store.pop(_job_key) print(f"[app][{_ts()}] all pages done", flush=True) _status_text.empty() _results = [ {"page": _i + 1, "bgr": _pages[_i], "dets": _all_dets[_i]} for _i in range(_n) ] # ── Build output files ──────────────────────────────────────────────── _status_text.text("Building output files…") _ann_pages = _annotate_pages(_results, _valid) st.session_state.ann_pages = _ann_pages st.session_state.ann_bytes = _pages_to_pdf_bytes(_ann_pages) st.session_state.csv_bytes = _build_csv_bytes(_results, st.session_state.legend_df) st.session_state.clean_bytes = _build_clean_bytes(_results, _ann_pages) st.session_state.results = _results st.session_state.valid = _valid st.session_state.last_run_summary = { "total_min": round((_time.time() - _t_pool) / 60), "pg_done": dict(_pg_done), "n": _n, "failed": False, } st.session_state.pending = None st.session_state.running = False st.rerun()