diff --git "a/app.py" "b/app.py" --- "a/app.py" +++ "b/app.py" @@ -0,0 +1,2534 @@ +""" +DeepPlan — Room Segmentation (SAM, prompt-point). SINGLE-FILE Gradio app. + +Runs the real DeepPlan wall extraction when the sibling modules are present, +otherwise an inline port of it. Deploy as `app.py`. + +Flow: + 1. Extract walls → wall_mask, by the wall_vectorizer_flask.py method ported + 1:1: preshrink → process_raster (invert, colour strip, ink threshold, + thickness, close) → vtracer trace of the stripped sheet → filter_wall_svg + (drop background colour, target colours, dashed runs, text blobs). The + filtered SVG is then rendered to the pixel mask this app consumes. Runs + entirely locally; nothing is uploaded. + 2. Measure every wall component's stroke and keep only the thick structural + ones. Sprinkler runs, electrical, plumbing, dimensions, text, symbols, + furniture and drafting lines fall out here, before anything downstream + can see them. + 3. Seal door openings and wall breaks, then label free space 4-connected. + A region is a room candidate only if it never touches the sheet edge and + its whole perimeter is wall — open, incomplete and exterior regions are + rejected outright. + 4. Prompt SAM per region: positive points at distance-transform maxima, + negative points on the wall ring and inside every adjacent region + (corridors, shafts, doorways, exterior), plus a tight box. One image + embedding for the sheet, one cheap mask-decoder call per room — not the + ~341 encoder passes SamAutomaticMaskGenerator's crop pyramid costs. + 5. Clip each mask at the wall and keep only the part connected to its seed, + so crossing a wall is structurally impossible. Rank SAM's three candidates + by IoU-with-region minus a leakage penalty. + 6. Validate the geometry — solidity, extent, vertex count, axis-alignment, + fragmentation — so only simple closed rectilinear rooms (square, rectangle, + L, T) survive. Merge over-segmented masks, drop contained duplicates. + 7. Show walls / prompts / colour-coded rooms / boundaries / per-room instances, + with live controls and PNG / SVG / JSON export. + +SAM checkpoint: CUDA if available else CPU. Auto-downloads vit_h (~2.4 GB) at boot +(disable warm-up with WARMUP_SAM=0; skip download with SAM_CHECKPOINT=/path.pth). + +Run: python sam_auto_app.py (PORT env, default 7861) +""" +from __future__ import annotations + +import json +import io +import math +import os +import re +import tempfile +import threading +import time +import xml.etree.ElementTree as ET +from contextlib import contextmanager +from typing import Any, Dict, Iterator, List, Optional, Tuple + +import cv2 +import numpy as np +from PIL import Image + +import gradio as gr + +# ───────────────────────────────────────────────────────────────────────────── +# REAL DeepPlan wall extraction — the exact `run_walls_only` the /extract-walls +# endpoint uses (pipeline.py + vector_walls + ocr + real stage1 DPI). This is the +# only path that is 1:1 with the frontend. Requires the sibling modules deployed +# beside this file: pipeline, constants, gpu_utils, ocr, sam_ops, vector_walls, +# room_validation, mask_to_polygon. Falls back to the inline port only if import +# fails. Force the port with USE_DEEPPLAN_WALLS=0. +# ───────────────────────────────────────────────────────────────────────────── +# Vector wall extraction from PDF. Optional (needs pymupdf); absent → the PDF tab +# reports why and the raster path is used. +try: + import pdf_walls # type: ignore + _HAS_PDF = pdf_walls.HAS_FITZ + _PDF_ERR = "" if _HAS_PDF else pdf_walls.FITZ_ERROR +except Exception as _pexc: + pdf_walls = None # type: ignore + _HAS_PDF = False + _PDF_ERR = f"{type(_pexc).__name__}: {_pexc}" +print(f"[pdf] vector wall extraction: {'available' if _HAS_PDF else 'OFF (' + _PDF_ERR + ')'}") + +_real_run_walls_only = None +_HAS_PIPELINE = False +if os.environ.get("USE_DEEPPLAN_WALLS", "1") != "0": + try: + from pipeline import run_walls_only as _real_run_walls_only # type: ignore + _HAS_PIPELINE = True + print("[walls] REAL DeepPlan pipeline.run_walls_only active (1:1 with frontend)") + except Exception as exc: + print(f"[walls] DeepPlan pipeline unavailable ({type(exc).__name__}: {exc}); " + f"using inline port (NOT 1:1). Deploy sibling modules for exact parity.") + +# ───────────────────────────────────────────────────────────────────────────── +# Inlined constants +# ──────────────────────────────────────────��────────────────────────────────── +SERVICE_VERSION = "6.0.0" +SEG_WALL_THICKEN_PX = int(os.environ.get("SEG_WALL_THICKEN_PX", "3")) +POLYGON_STRAIGHTEN_DEG = float(os.environ.get("POLYGON_STRAIGHTEN_DEG", "10.0")) + +# Colored-architecture thresholds (HSV S/V are 0-255, OpenCV convention). +SAT_MEP_MIN = int(os.environ.get("SAT_MEP_MIN", "150")) +SAT_WALL_MIN = int(os.environ.get("SAT_WALL_MIN", "40")) +WALL_COLOR_V_MIN = int(os.environ.get("WALL_COLOR_V_MIN", "40")) +WALL_COLOR_V_MAX = int(os.environ.get("WALL_COLOR_V_MAX", "245")) +NEAR_WHITE_S_MAX = int(os.environ.get("NEAR_WHITE_S_MAX", "30")) +NEAR_WHITE_V_MIN = int(os.environ.get("NEAR_WHITE_V_MIN", "230")) +COLORED_ARCH_RATIO_MIN = float(os.environ.get("COLORED_ARCH_RATIO_MIN", "0.10")) + +# OCR text-erase (Stage 4 step 3). +MAX_PROCESSING_DIMENSION = int(os.environ.get("MAX_PROCESSING_DIMENSION", "4000")) +# Wall extraction runs at FULL resolution (frontend parity) — downscaling before +# extraction welds fine line gaps on dense sheets → blob walls. Cap only for memory. +EXTRACT_MAX_DIM = int(os.environ.get("EXTRACT_MAX_DIM", "4000")) +OCR_ERASE_CONF_FLOOR = float(os.environ.get("OCR_ERASE_CONF_FLOOR", "0.20")) +OCR_ERASE_DILATION_PX = int(os.environ.get("OCR_ERASE_DILATION_PX", "4")) + +SAM_MODEL_TYPE = os.environ.get("SAM_MODEL_TYPE", "vit_h") +SAM_CHECKPOINT_NAME = os.environ.get("SAM_CHECKPOINT_NAME", "sam_vit_h_4b8939.pth") +SAM_CHECKPOINT_URL = os.environ.get( + "SAM_CHECKPOINT_URL", + "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth", +) +SAM_CHECKPOINT_PATH = os.environ.get("SAM_CHECKPOINT", "") + +# Per-room overlay palette (RGB). +ROOM_COLORS = ( + (255, 99, 132), (54, 162, 235), (255, 206, 86), (75, 192, 192), + (153, 102, 255), (255, 159, 64), (231, 233, 237), (199, 92, 122), + (132, 220, 198), (255, 140, 105), (180, 200, 255), (140, 220, 140), + (255, 180, 80), (200, 150, 255), (100, 220, 220), (240, 130, 200), +) + +# SamAutomaticMaskGenerator defaults requested for this app. +DEFAULTS = dict( + points_per_side=64, + pred_iou_thresh=0.90, + stability_score_thresh=0.95, + crop_n_layers=4, + crop_overlap_ratio=512 / 1500, + crop_n_points_downscale_factor=2, + min_mask_region_area=300, + points_per_batch=int(os.environ.get("SAM_POINTS_PER_BATCH", "64")), +) + +_EXPORT_DIR = os.path.join(tempfile.gettempdir(), "sam_auto_exports") +os.makedirs(_EXPORT_DIR, exist_ok=True) + + +# ═════════════════════════════════════════════════════════════════════════════ +# WALL EXTRACTION — 1:1 port of wall_vectorizer_flask.py. +# +# Every constant, threshold and step below is the Flask module verbatim. The +# pipeline it defines is: +# +# preshrink Pillow downscale to MAX_PIXELS (off by default). +# process_raster polarity normalise, strip chromatic MEP, ink threshold, +# thickness filter, gap close. Returns the colour-stripped +# sheet and a black-walls-on-white raster. +# vtrace_svg vtracer traces the STRIPPED sheet (not the wall raster — +# this is what the Flask app feeds it) to an SVG of the whole +# drawing. +# filter_wall_svg keep the wall vectors: drop the dominant/background colour, +# any DROP_COLORS, dashed runs (opt-in) and small compact text +# blobs. Emits the wall-only SVG. +# +# The Flask app's deliverable is that SVG. This app consumes a pixel mask, so +# `_wall_svg_to_mask` below renders the filtered SVG and thresholds it. That +# renderer is the only thing here that is NOT in the Flask module — it is the +# adapter, kept separate so the ported method stays untouched. +# +# Runs entirely on this host: no API, no credits, no blueprint leaving the +# machine. VTRACE_WALLS=0 disables wall extraction entirely. +# ═════════════════════════════════════════════════════════════════════════════ +try: + import vtracer # type: ignore + _HAS_VTRACER = True + _VTRACER_ERR = "" +except Exception as _vexc: # pragma: no cover + vtracer = None # type: ignore + _HAS_VTRACER = False + _VTRACER_ERR = f"{type(_vexc).__name__}: {_vexc}" + +try: + from svgpathtools import parse_path # type: ignore + _HAS_SVGPATHTOOLS = True + _SVGPT_ERR = "" +except Exception as _sexc: # pragma: no cover + parse_path = None # type: ignore + _HAS_SVGPATHTOOLS = False + _SVGPT_ERR = f"{type(_sexc).__name__}: {_sexc}" + +print(f"[vtrace] local vector wall extraction: " + f"{'available' if _HAS_VTRACER else 'OFF (' + _VTRACER_ERR + ')'}" + f"{'' if _HAS_SVGPATHTOOLS else ' · svgpathtools missing (' + _SVGPT_ERR + ')'}") + +VTRACE_WALLS = os.environ.get("VTRACE_WALLS", "1") == "1" + +# ───────────────────────────────────────────────────────────────────────────── +# Config — wall_vectorizer_flask.py, verbatim +# ───────────────────────────────────────────────────────────────────────────── +# vtracer trace params (see visioncortex/vtracer Python Config). +VTRACER_CLUSTERING = os.environ.get("VTRACER_CLUSTERING", "color-cluster") # color-cluster|bw|watershed +VTRACER_MODE = os.environ.get("VTRACER_MODE", "polygon") # pixel|polygon|spline +VTRACER_HIERARCHICAL = os.environ.get("VTRACER_HIERARCHICAL", "stacked") # stacked|cutout +VTRACER_FILTER_SPECKLE = int(os.environ.get("VTRACER_FILTER_SPECKLE", "1")) # low = keep more thin wall detail +VTRACER_COLOR_PRECISION = int(os.environ.get("VTRACER_COLOR_PRECISION", "6")) +VTRACER_CORNER_THRESHOLD = int(os.environ.get("VTRACER_CORNER_THRESHOLD", "60")) +VTRACER_PATH_PRECISION = int(os.environ.get("VTRACER_PATH_PRECISION", "5")) + +# OpenCV colour strip (raster, BEFORE tracing): whiten chromatic (red/green/blue MEP) +# pixels so only grayscale walls/text reach vtracer. HSV S/V 0-255. +STRIP_SAT_MIN = int(os.environ.get("STRIP_SAT_MIN", "40")) # saturation >= this = colored -> whitened +STRIP_DILATE_PX = int(os.environ.get("STRIP_DILATE_PX", "1")) # grow colored mask to catch anti-aliased edges + +# Solid-wall isolation on the stripped raster (OpenCV): +WALL_INK_MAX = int(os.environ.get("WALL_INK_MAX", "150")) # keep pixels darker than this; faint gray dims/grid dropped +WALL_MIN_STROKE = int(os.environ.get("WALL_MIN_STROKE", "0")) # keep only CCs with a >= this-px-thick core (0 = off) +WALL_CLOSE_PX = int(os.environ.get("WALL_CLOSE_PX", "2")) # morph-close to seal gaps -> continuous solid walls + +# Grayscale-keep filter: walls are GRAYSCALE ink (black..gray). Keep near-neutral, +# dark-enough paths; reject colored MEP (saturated). +GRAY_MAX = int(os.environ.get("GRAY_MAX", "250")) # brightest channel allowed (higher = keep lighter gray) +NEUTRAL_TOL = int(os.environ.get("NEUTRAL_TOL", "20")) # max channel spread = achromatic; ANY red/green/blue tint above this = dropped +# Text removal: a wall subpath is long; a text glyph is short. Drop grayscale +# subpaths whose bbox diagonal is below this. +WALL_MIN_LEN = float(os.environ.get("WALL_MIN_LEN", "30")) + +# ── Dashed/dotted removal (keep ALL solid lines) ───────────────────────────── +# vtracer traces each dash/dot as its own short subpath. A dashed line = a +# collinear run of short marks with regular gaps. Detect runs -> drop; keep +# everything else (all solid lines, any length; isolated short marks stay). +# DEFAULT OFF: vtracer color-cluster fragments a SOLID wall into many small +# same-gray adjacent polygons — a collinear run — which the dash detector +# wrongly kills. Since a fragmented solid wall is indistinguishable from a real +# dashed line, keep dash removal opt-in (DASH_ENABLE=1) to never lose solids. +DASH_ENABLE = os.environ.get("DASH_ENABLE", "0") == "1" # 1 = also strip dashed runs (risks fragmented solids) +DASH_MAX_LEN = float(os.environ.get("DASH_MAX_LEN", "22")) # a dash/dot mark's max bbox diagonal px +DASH_SNAP = float(os.environ.get("DASH_SNAP", "4")) # collinear row/col bucket px +DASH_GAP_MIN = float(os.environ.get("DASH_GAP_MIN", "4")) # min gap between consecutive dashes px (solid fragments touch, gap~0 -> NOT dashed) +DASH_GAP_MAX = float(os.environ.get("DASH_GAP_MAX", "30")) # max gap between consecutive dashes px +DASH_MIN_RUN = int(os.environ.get("DASH_MIN_RUN", "5")) # >= this many collinear marks = dashed line + +# ── Text-blob removal (small + compact marks) ──────────────────────────────── +# Text glyphs / numbers / symbols = small AND near-square (low aspect). Solid +# wall lines = elongated (high aspect) -> kept at ANY length. Drop only small +# compact blobs. +TEXT_MAX_DIAG = float(os.environ.get("TEXT_MAX_DIAG", "28")) # only marks this small can be text px +TEXT_MAX_ASPECT = float(os.environ.get("TEXT_MAX_ASPECT", "2.4")) # long/short bbox below this = compact = text + +# ── Colour-based removal (major/background colour + explicit dashed colour) ─── +# The dominant fill by area = blueprint background. Drop subpaths whose fill is +# that background colour (removes near-white filler polygons). Also drop any +# fill listed in DROP_COLORS (e.g. the dashed-line colour), within DROP_TOL. +DROP_BG = os.environ.get("DROP_BG", "1") == "1" # drop subpaths matching the dominant background colour +BG_TOL = int(os.environ.get("BG_TOL", "6")) # channel tolerance vs background colour +DROP_COLORS = os.environ.get("DROP_COLORS", "") # comma hex list to drop, e.g. "#000000,#010101" (dashed) +DROP_TOL = int(os.environ.get("DROP_TOL", "10")) # channel tolerance vs each DROP_COLORS entry + +# vtracer is local (no API pixel cap), but tracing cost grows with pixels; pre-shrink +# huge sheets for speed. Raise for finer walls. +MAX_PIXELS = int(os.environ.get("MAX_PIXELS", "6000000")) +PRESHRINK_ENABLE = os.environ.get("PRESHRINK_ENABLE", "0") == "1" # 0 = never downscale (full-res trace) +AUTO_INVERT = os.environ.get("AUTO_INVERT", "1") == "1" # invert dark-theme plots -> black ink on white bg +INVERT_THRESH = int(os.environ.get("INVERT_THRESH", "128")) # mean luminance below this = dark bg -> invert + +_SVG_NS = "http://www.w3.org/2000/svg" + + +def preshrink(image_bytes: bytes, log: List[str]) -> bytes: + """Downscale (Pillow only, not for extraction) so pixel count <= MAX_PIXELS, for + trace speed. Returns PNG bytes.""" + try: + im = Image.open(io.BytesIO(image_bytes)) + im.load() + except Exception as exc: + log.append(f"preshrink: could not open image ({exc}) — using raw bytes") + return image_bytes + w, h = im.size + if PRESHRINK_ENABLE and w * h > MAX_PIXELS: + s = math.sqrt(MAX_PIXELS / float(w * h)) + nw, nh = max(1, int(w * s)), max(1, int(h * s)) + im = im.resize((nw, nh), Image.LANCZOS) + log.append(f"preshrink: {w}x{h} ({w*h:,}px) -> {nw}x{nh} ({nw*nh:,}px)") + else: + log.append(f"preshrink: {w}x{h} ({w*h:,}px) within limit — unchanged") + im = im.convert("RGB") + buf = io.BytesIO() + im.save(buf, format="PNG") + return buf.getvalue() + + +def _encode(img: np.ndarray) -> bytes: + ok, out = cv2.imencode(".png", img) + return out.tobytes() if ok else b"" + + +def process_raster(png_bytes: bytes, log: List[str]) -> Tuple[bytes, bytes]: + """OpenCV wall isolation. Returns (stripped_png, wall_png = black solid walls on + white), the latter fed to vtracer. + + 1. Strip colour — whiten chromatic (S >= STRIP_SAT_MIN) MEP linework. + 2. Ink threshold — keep only DARK pixels (< WALL_INK_MAX); faint gray dimension + / grid / leader lines drop out, structural walls stay. + 3. Thickness — optional: keep only connected components with a thick core + (>= WALL_MIN_STROKE px), so thin annotation lines that were still dark go. + 4. Close — seal small gaps so walls are continuous solids. + """ + bgr = cv2.imdecode(np.frombuffer(png_bytes, np.uint8), cv2.IMREAD_COLOR) + if bgr is None: + log.append("process: could not decode — skipped") + return png_bytes, png_bytes + + # 0. normalize polarity — dark-theme CAD plots (light ink on dark bg) -> invert + # so downstream always sees black ink on white background. + if AUTO_INVERT: + mean_lum = float(cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY).mean()) + if mean_lum < INVERT_THRESH: + bgr = 255 - bgr + log.append(f"invert: dark bg detected (mean lum {mean_lum:.0f} < {INVERT_THRESH}) -> inverted") + + # 1. strip colour + hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV) + colored = (hsv[:, :, 1] >= STRIP_SAT_MIN).astype(np.uint8) + if STRIP_DILATE_PX > 0: + k = cv2.getStructuringElement( + cv2.MORPH_ELLIPSE, (2 * STRIP_DILATE_PX + 1, 2 * STRIP_DILATE_PX + 1)) + colored = cv2.dilate(colored, k, iterations=1) + bgr[colored > 0] = (255, 255, 255) + log.append(f"strip: whitened {int(np.count_nonzero(colored)):,} colored px (S>={STRIP_SAT_MIN})") + stripped_png = _encode(bgr) + + # 2. dark ink only + gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY) + ink = (gray < WALL_INK_MAX).astype(np.uint8) * 255 + log.append(f"ink: {int(np.count_nonzero(ink)):,} dark px (< {WALL_INK_MAX})") + + # 3. thickness filter — keep components with a thick core + if WALL_MIN_STROKE > 0: + dist = cv2.distanceTransform(ink, cv2.DIST_L2, 5) + _n, lbl = cv2.connectedComponents(ink, connectivity=8) + thick = np.unique(lbl[dist >= WALL_MIN_STROKE / 2.0]) + thick = thick[thick != 0] + ink = np.isin(lbl, thick).astype(np.uint8) * 255 + log.append(f"thickness: kept {len(thick)} thick CCs (>= {WALL_MIN_STROKE}px core)") + + # 4. close gaps -> continuous solids + if WALL_CLOSE_PX > 0: + k = cv2.getStructuringElement( + cv2.MORPH_ELLIPSE, (2 * WALL_CLOSE_PX + 1, 2 * WALL_CLOSE_PX + 1)) + ink = cv2.morphologyEx(ink, cv2.MORPH_CLOSE, k) + + wall_bow = np.where(ink > 0, 0, 255).astype(np.uint8) # black walls on white + return stripped_png, _encode(wall_bow) + + +# ───────────────────────────────────────────────────────────────────────────── +# vtracer raster->vector (local) +# ───────────────────────��───────────────────────────────────────────────────── +def vtrace_svg(png_bytes: bytes, log: List[str]) -> str: + """Trace raster -> SVG string with vtracer, locally. Supports both bindings: the + newer Config API and the older convert_raw_image_to_svg function API.""" + if hasattr(vtracer, "Config"): + cfg = vtracer.Config( + clustering=VTRACER_CLUSTERING, hierarchical=VTRACER_HIERARCHICAL, + mode=VTRACER_MODE, filter_speckle=VTRACER_FILTER_SPECKLE, + color_precision=VTRACER_COLOR_PRECISION, + corner_threshold=VTRACER_CORNER_THRESHOLD, path_precision=VTRACER_PATH_PRECISION, + ) + svg = cfg.convert_bytes(png_bytes) + api = "Config" + else: + colormode = "binary" if VTRACER_CLUSTERING == "bw" else "color" + svg = vtracer.convert_raw_image_to_svg( + png_bytes, img_format="png", colormode=colormode, + hierarchical=VTRACER_HIERARCHICAL, mode=VTRACER_MODE, + filter_speckle=VTRACER_FILTER_SPECKLE, color_precision=VTRACER_COLOR_PRECISION, + corner_threshold=VTRACER_CORNER_THRESHOLD, path_precision=VTRACER_PATH_PRECISION, + ) + api = "convert_raw_image_to_svg" + log.append(f"vtracer[{api}]: {VTRACER_CLUSTERING}/{VTRACER_MODE}, " + f"speckle={VTRACER_FILTER_SPECKLE} -> {len(svg):,} bytes SVG") + return svg + + +# ───────────────────────────────────────────────────────────────────────────── +# SVG parse + wall-vector filter (pure geometry + colour) +# ───────────────────────────────────────────────────────────────────────────── +def _local(tag: str) -> str: + return tag.rsplit("}", 1)[-1] + + +def _iter_geometry(root: ET.Element): + """Yield (element, 'd'-string) for every drawable path/shape.""" + for el in root.iter(): + t = _local(el.tag) + if t == "path" and el.get("d"): + yield el, el.get("d") + elif t == "polygon" and el.get("points"): + yield el, "M " + el.get("points").strip() + " Z" + elif t == "polyline" and el.get("points"): + yield el, "M " + el.get("points").strip() + elif t == "line": + yield el, (f"M {el.get('x1','0')},{el.get('y1','0')} " + f"L {el.get('x2','0')},{el.get('y2','0')}") + elif t == "rect": + x = float(el.get("x", 0)); y = float(el.get("y", 0)) + w = float(el.get("width", 0)); h = float(el.get("height", 0)) + yield el, f"M {x},{y} L {x+w},{y} L {x+w},{y+h} L {x},{y+h} Z" + + +_NAMED = {"black": (0, 0, 0), "white": (255, 255, 255), "red": (255, 0, 0), + "green": (0, 128, 0), "blue": (0, 0, 255), "none": None} + + +def _parse_color(val: str): + if not val: + return None + v = val.strip().lower() + if v in _NAMED: + return _NAMED[v] + if v.startswith("#"): + h = v[1:] + if len(h) == 3: + h = "".join(c * 2 for c in h) + if len(h) == 6: + try: + return int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16) + except ValueError: + return None + if v.startswith("rgb"): + try: + nums = v[v.index("(") + 1:v.index(")")].split(",") + return tuple(int(float(n.strip().rstrip("%"))) for n in nums[:3]) + except Exception: + return None + return None + + +def _paint_rgb(el: ET.Element): + fill = el.get("fill") + stroke = el.get("stroke") + for decl in (el.get("style") or "").split(";"): + if ":" in decl: + k, val = decl.split(":", 1) + if k.strip() == "fill": + fill = val.strip() + elif k.strip() == "stroke": + stroke = val.strip() + if fill and fill.strip().lower() != "none": + return _parse_color(fill) + if stroke and stroke.strip().lower() != "none": + return _parse_color(stroke) + return _parse_color(fill) + + +def is_wall_ink(el: ET.Element) -> bool: + """Grayscale wall ink: near-neutral (channel spread <= NEUTRAL_TOL) AND dark + enough (brightest channel <= GRAY_MAX). Keeps black..gray walls, drops saturated + colored MEP. Untinted paths default to black in SVG.""" + rgb = _paint_rgb(el) + if rgb is None: + return True + if max(rgb) - min(rgb) > NEUTRAL_TOL: # colored (saturated) -> not a wall + return False + return max(rgb) <= GRAY_MAX + + +def _attrs_str(el: ET.Element, d: str) -> str: + """Serialize an element back to a , preserving its original paint so it + renders exactly as vtracer drew it (grayscale, correct fill-rule).""" + out = [f'd="{d}"'] + for k in ("fill", "fill-rule", "stroke", "stroke-width", "opacity", "transform", "style"): + v = el.get(k) + if v: + out.append(f'{k}="{v}"') + return " ".join(out) + + +def _detect_dashed(recs: List[dict]) -> set: + """Flag indices of subpaths that belong to a dashed/dotted line: collinear runs + of >=DASH_MIN_RUN short marks spaced DASH_GAP_MIN..MAX apart. Returns dashed set.""" + cand = [i for i, r in enumerate(recs) if r["diag"] <= DASH_MAX_LEN] + buckets = {} + for i in cand: + r = recs[i] + o = r["orient"] + fixed = r["cy"] if o == "h" else r["cx"] # across-axis coord + buckets.setdefault((o, round(fixed / DASH_SNAP)), []).append(i) + dashed = set() + for (o, _k), idxs in buckets.items(): + if len(idxs) < DASH_MIN_RUN: + continue + # sort along axis by mark start; extent along axis = w (h-orient) or h (v-orient) + def along(i): + r = recs[i] + c = r["cx"] if o == "h" else r["cy"] + ext = (r["w"] if o == "h" else r["h"]) / 2.0 + return c - ext, c + ext + idxs.sort(key=lambda i: along(i)[0]) + run = [idxs[0]] + _, prev_end = along(idxs[0]) + for i in idxs[1:]: + s, e = along(i) + gap = s - prev_end + if DASH_GAP_MIN <= gap <= DASH_GAP_MAX: + run.append(i) + else: + if len(run) >= DASH_MIN_RUN: + dashed.update(run) + run = [i] + prev_end = max(prev_end, e) + if len(run) >= DASH_MIN_RUN: + dashed.update(run) + return dashed + + +def filter_wall_svg(svg_text: str, log: List[str]) -> Tuple[str, int, int]: + """Parse the traced SVG, keep ALL solid lines, drop only dashed/dotted lines. + Returns (wall_svg, kept, total).""" + try: + root = ET.fromstring(svg_text) + except ET.ParseError as exc: + raise RuntimeError(f"SVG parse failed: {exc}") + width = root.get("width", "") + height = root.get("height", "") + view_box = root.get("viewBox") or root.get("viewbox") or "" + + # Pass 1: collect EVERY subpath with bbox geometry. No color/darkness erase — + # only dashed + text get removed below (Step 2). MEP colour already stripped + # from the raster upstream. + recs: List[dict] = [] # each: el, d, cx, cy, w, h, diag, orient, rgb + total = 0 + color_area = {} # rgb -> total bbox area (for dominant/background) + for el, d in _iter_geometry(root): + try: + p = parse_path(d) + except Exception: + continue + rgb = _paint_rgb(el) + for sub in p.continuous_subpaths(): + total += 1 + try: + xmin, xmax, ymin, ymax = sub.bbox() + except Exception: + continue + w = float(xmax - xmin) + h = float(ymax - ymin) + aspect = max(w, h) / max(min(w, h), 1e-6) + if rgb is not None: + color_area[rgb] = color_area.get(rgb, 0.0) + w * h + recs.append({ + "el": el, "d": sub.d(), + "cx": float(xmin + xmax) / 2.0, "cy": float(ymin + ymax) / 2.0, + "w": w, "h": h, "diag": math.hypot(w, h), "aspect": aspect, + "orient": "h" if w >= h else "v", "rgb": rgb, + }) + + # Dominant colour by area = blueprint background. Log the histogram so the + # dashed-line colour is identifiable for DROP_COLORS. + ranked = sorted(color_area.items(), key=lambda kv: -kv[1]) + bg_rgb = ranked[0][0] if ranked else None + top = ", ".join(f"#{r:02X}{g:02X}{b:02X}({a:.0f})" for (r, g, b), a in ranked[:8]) + log.append(f"colors: dominant/bg={('#%02X%02X%02X' % bg_rgb) if bg_rgb else 'n/a'}; top by area: {top}") + + drop_targets = [c for c in (_parse_color(x) for x in DROP_COLORS.split(",")) if c] + + def _near(a, b, tol): + return a is not None and b is not None and all(abs(a[i] - b[i]) <= tol for i in range(3)) + + # Pass 2: optionally flag dashed/dotted runs. OFF by default so color-cluster + # fragmented solid walls are never mistaken for dashes. + dashed = _detect_dashed(recs) if DASH_ENABLE else set() + + # Re-emit kept subpaths grouped by their original element (preserve paint/fill-rule). + by_el = {} + order = [] + kept_subs = 0 + dropped_text = 0 + dropped_bg = 0 + dropped_color = 0 + for i, r in enumerate(recs): + if i in dashed: + continue + # Background colour -> drop (near-white filler polygons). + if DROP_BG and _near(r["rgb"], bg_rgb, BG_TOL): + dropped_bg += 1 + continue + # Explicit dashed / target colours -> drop. + if drop_targets and any(_near(r["rgb"], t, DROP_TOL) for t in drop_targets): + dropped_color += 1 + continue + # Text blob: small AND compact (near-square). Elongated lines survive at any length. + if r["diag"] <= TEXT_MAX_DIAG and r["aspect"] <= TEXT_MAX_ASPECT: + dropped_text += 1 + continue + key = id(r["el"]) + if key not in by_el: + by_el[key] = (r["el"], []) + order.append(key) + by_el[key][1].append(r["d"]) + kept_subs += 1 + kept = [f" " for k in order] + + hdr = [f'") + svg = "".join(hdr) + "\n" + "\n".join(kept) + "\n" + log.append(f"filter: kept {kept_subs}/{total} subpaths (dropped {dropped_bg} bg, " + f"{dropped_color} target-color, {len(dashed)} dashed, {dropped_text} text blobs; " + f"bg tol={BG_TOL}, drop_colors={DROP_COLORS or 'none'}/tol={DROP_TOL})") + return svg, kept_subs, total + + +# ═════════════════════════════════════════════════════════════════════════════ +# ADAPTER — wall SVG -> pixel mask. NOT part of the ported method. +# +# The Flask app ships the SVG and lets a browser draw it; this app needs pixels, +# so the filtered SVG is rendered here and thresholded. Rendering is done in +# document order with each path's own fill, because vtracer's `stacked` output +# means a later shape is meant to cover an earlier one — a union of the same +# polygons would flood the sheet. Subpaths of one element are filled even-odd so +# a hollow outline stays hollow. +# ═════════════════════════════════════════════════════════════════════════════ +_NUM_RE = r"[-+]?\d*\.?\d+(?:[eE][-+]?\d+)?" +_TOKEN_RE = re.compile(rf"([MLZmlz])|({_NUM_RE})") +_TRANSLATE_RE = re.compile(r"translate\(\s*([-\d.eE+]+)[ ,]+([-\d.eE+]+)\s*\)", re.I) + + +def _svg_subpaths(d: str) -> List[np.ndarray]: + """'d' string -> point arrays, one per subpath. vtracer is asked for polygon + mode, so only M/L/Z appear.""" + subs: List[np.ndarray] = [] + cur: List[Tuple[float, float]] = [] + nums: List[float] = [] + cmd: Optional[str] = None + + def flush() -> None: + if len(cur) >= 3: + subs.append(np.array(cur, dtype=np.float64)) + + for c, n in _TOKEN_RE.findall(d or ""): + if c: + if c in "Zz": + flush() + cur = [] + else: + if c in "Mm" and cur: + flush() + cur = [] + cmd = c + nums = [] + continue + nums.append(float(n)) + if len(nums) == 2: + x, y = nums + nums = [] + if cmd in ("m", "l") and cur: # relative + x += cur[-1][0] + y += cur[-1][1] + cur.append((x, y)) + flush() + return subs + + +def _wall_svg_to_mask(wall_svg: str, shape: Tuple[int, int], + log: List[str]) -> np.ndarray: + """Render the wall SVG and threshold it into a binary wall mask.""" + h, w = shape + canvas = np.full((h, w), 255, np.uint8) + root = ET.fromstring(wall_svg) + painted = 0 + for el in root.iter(): + if _local(el.tag) != "path" or not el.get("d"): + continue + subs = _svg_subpaths(el.get("d") or "") + if not subs: + continue + t = _TRANSLATE_RE.search(el.get("transform") or "") + if t: + off = np.array([float(t.group(1)), float(t.group(2))], dtype=np.float64) + subs = [s + off for s in subs] + rgb = _paint_rgb(el) + # An untinted path defaults to black fill in SVG, i.e. ink. + lum = (0 if rgb is None else + int(round(0.299 * rgb[0] + 0.587 * rgb[1] + 0.114 * rgb[2]))) + pts_all = np.concatenate(subs, axis=0) + x0 = max(0, int(np.floor(pts_all[:, 0].min()))) + y0 = max(0, int(np.floor(pts_all[:, 1].min()))) + x1 = min(w, int(np.ceil(pts_all[:, 0].max())) + 1) + y1 = min(h, int(np.ceil(pts_all[:, 1].max())) + 1) + if x1 <= x0 or y1 <= y0: + continue + local = np.zeros((y1 - y0, x1 - x0), np.uint8) + one = np.empty_like(local) + for poly in subs: + one[:] = 0 + cv2.fillPoly(one, [np.round(poly - (x0, y0)).astype(np.int32)], 1) + local ^= one # even-odd + canvas[y0:y1, x0:x1][local > 0] = lum + painted += 1 + mask = (canvas < WALL_INK_MAX).astype(np.uint8) * 255 + log.append(f"raster: rendered {painted} kept paths -> " + f"{int(np.count_nonzero(mask)):,} wall px " + f"({100.0 * np.count_nonzero(mask) / mask.size:.1f}%, ink < {WALL_INK_MAX})") + return mask + + +def extract_walls_via_vtracer(bgr: np.ndarray, log: List[str]) -> Optional[np.ndarray]: + """The wall_vectorizer_flask.py /api/extract flow, end to end, on a BGR image. + + preshrink -> process_raster -> vtrace_svg(stripped) -> filter_wall_svg, then + the adapter renders the wall SVG into the pixel mask this app consumes. The + filtered SVG is also saved for export. Returns None if no walls survive.""" + if not _HAS_VTRACER: + log.append(f"[vtrace] vtracer not installed ({_VTRACER_ERR}) — pip install vtracer") + return None + if not _HAS_SVGPATHTOOLS: + log.append(f"[vtrace] svgpathtools not installed ({_SVGPT_ERR}) — pip install svgpathtools") + return None + try: + png = preshrink(_encode(bgr), log) + stripped, _wall_png = process_raster(png, log) + full_svg = vtrace_svg(stripped, log) + wall_svg, kept, total = filter_wall_svg(full_svg, log) + if kept == 0: + log.append("[vtrace] filter kept no subpaths") + return None + + mask = _wall_svg_to_mask(wall_svg, bgr.shape[:2], log) + if not np.any(mask): + log.append("[vtrace] rendered wall mask is empty") + return None + + svg_path = os.path.join(_EXPORT_DIR, f"walls_vector_{int(time.time())}.svg") + try: + with open(svg_path, "w", encoding="utf-8") as fh: + fh.write(wall_svg) + log.append(f"[vtrace] walls.svg saved {svg_path}") + except Exception as exc: + log.append(f"[vtrace] svg save failed ({type(exc).__name__}: {exc})") + return mask + except Exception as exc: + log.append(f"[vtrace] failed ({type(exc).__name__}: {exc})") + return None + + +def extract_walls_primary(bgr: np.ndarray, log: List[str], dpi: int = 0, + scale_denominator: int = 100) -> Tuple[Optional[np.ndarray], str]: + """Wall source for the raster path: the ported wall_vectorizer_flask.py + method, which is the only one. Nothing is uploaded anywhere. `dpi` and + `scale_denominator` are accepted for call-site compatibility; the ported + method is defined in pixels and does not use them.""" + if not VTRACE_WALLS: + log.append("[walls] VTRACE_WALLS=0 — wall extraction disabled") + return None, "disabled" + mask = extract_walls_via_vtracer(bgr, log) + if mask is None: + return None, "vtracer (failed)" + return mask, "wall_vectorizer_flask vector walls" + + +# ───────────────────────────────────────────────────────────────────────────── +# Inlined SAM loader — torch/CUDA detection + lazy checkpoint + session lock +# ───────────────────────────────────────────────────────────────────────────── +try: + import torch # type: ignore + _HAS_TORCH = True + _CUDA = torch.cuda.is_available() + _DEVICE = "cuda" if _CUDA else "cpu" + _GPU_NAME = torch.cuda.get_device_name(0) if _CUDA else "CPU" +except ImportError: + torch = None # type: ignore + _HAS_TORCH = False + _CUDA = False + _DEVICE = "cpu" + _GPU_NAME = "CPU (no torch)" + +try: + from segment_anything import ( # type: ignore + sam_model_registry, SamPredictor, SamAutomaticMaskGenerator, + ) + _HAS_SAM = True +except ImportError: + sam_model_registry = None # type: ignore + SamPredictor = None # type: ignore + SamAutomaticMaskGenerator = None # type: ignore + _HAS_SAM = False + +_sam_predictor: Optional[Any] = None +_sam_lock = threading.RLock() + + +def get_sam_predictor() -> Optional[Any]: + """Lazy-load SAM predictor. Downloads checkpoint if absent. None if unavailable.""" + global _sam_predictor + if _sam_predictor is not None: + return _sam_predictor + if not _HAS_SAM or not _HAS_TORCH: + return None + + cache_dir = os.path.join(tempfile.gettempdir(), "sam_cache") + os.makedirs(cache_dir, exist_ok=True) + + ckpt_path = SAM_CHECKPOINT_PATH.strip() + if not ckpt_path or not os.path.isfile(ckpt_path): + ckpt_path = os.path.join(cache_dir, SAM_CHECKPOINT_NAME) + if not os.path.isfile(ckpt_path): + try: + import urllib.request + print(f"[SAM] Downloading {SAM_CHECKPOINT_NAME}...") + urllib.request.urlretrieve(SAM_CHECKPOINT_URL, ckpt_path) + except Exception as exc: + print(f"[SAM] Download failed: {exc}") + return None + try: + sam = sam_model_registry[SAM_MODEL_TYPE](checkpoint=ckpt_path) + sam.to(device=_DEVICE) + sam.eval() + _sam_predictor = SamPredictor(sam) + print(f"[SAM] Ready on {_DEVICE}") + except Exception as exc: + print(f"[SAM] Load failed: {exc}") + return None + return _sam_predictor + + +@contextmanager +def sam_session() -> Iterator[Optional[Any]]: + with _sam_lock: + yield get_sam_predictor() + + +# ───────────────────────────────────────────────────────────────────────────── +# Inlined OCR text-erase — EasyOCR singleton + raw hits + bbox paint. +# ───────────────────────────────────────────────────────────────────────────── +try: + import easyocr # type: ignore + _HAS_EASYOCR = True +except ImportError: + easyocr = None # type: ignore + _HAS_EASYOCR = False + +_ocr_reader: Optional[Any] = None + + +def get_ocr_reader() -> Optional[Any]: + global _ocr_reader + if _ocr_reader is None and _HAS_EASYOCR and easyocr is not None: + _ocr_reader = easyocr.Reader(["en"], gpu=_CUDA, verbose=False) + return _ocr_reader + + +def run_ocr_raw(bgr_source: np.ndarray) -> List[Dict[str, Any]]: + """EasyOCR once on bgr_source → hits at original-image coords. Downscale to + MAX_PROCESSING_DIMENSION, grayscale, CLAHE, drop conf < OCR_ERASE_CONF_FLOOR.""" + if not _HAS_EASYOCR: + return [] + reader = get_ocr_reader() + if reader is None: + return [] + h, w = bgr_source.shape[:2] + scale_factor = 1.0 + if max(h, w) > MAX_PROCESSING_DIMENSION: + scale_factor = MAX_PROCESSING_DIMENSION / max(h, w) + ocr_input = cv2.resize(bgr_source, (int(w * scale_factor), int(h * scale_factor)), + interpolation=cv2.INTER_AREA) + else: + ocr_input = bgr_source + enhanced = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)).apply( + cv2.cvtColor(ocr_input, cv2.COLOR_BGR2GRAY)) + try: + results = reader.readtext(enhanced, detail=1, paragraph=False) + except Exception as exc: + print(f"[Stage 4 OCR-erase] failed: {exc}") + return [] + hits: List[Dict[str, Any]] = [] + for bbox, text, conf in results: + if conf < OCR_ERASE_CONF_FLOOR: + continue + pts = (np.array(bbox, dtype=np.float32) / scale_factor).astype(np.int32) + hits.append({"bbox": pts.tolist(), "conf": float(conf)}) + return hits + + +def _ocr_erase_text_on_wall(wall_mask: np.ndarray, raw_hits: List[Dict[str, Any]], + dilation_px: int = OCR_ERASE_DILATION_PX) -> int: + """Paint dilated axis-aligned bbox of every OCR hit to 0 on wall_mask. In-place.""" + erased = 0 + for hit in raw_hits: + pts = np.array(hit["bbox"], dtype=np.int32) + x0 = max(0, int(pts[:, 0].min()) - dilation_px) + y0 = max(0, int(pts[:, 1].min()) - dilation_px) + x1 = min(wall_mask.shape[1], int(pts[:, 0].max()) + dilation_px) + y1 = min(wall_mask.shape[0], int(pts[:, 1].max()) + dilation_px) + if x1 > x0 and y1 > y0: + wall_mask[y0:y1, x0:x1] = 0 + erased += 1 + return erased + + +# ═════════════════════════════════════════════════════════════════════════════ +# HEURISTIC WALL EXTRACTION — self-contained port of DeepPlan Stages 2-4. +# (pipeline.py stage2_crop_drawing / stage3_strip_colors / stage4_extract_walls, +# steps 1-9. OCR text-erase and vector_walls refinement are omitted — the former +# needs easyocr, the latter is a separate deterministic refinement module.) +# ═════════════════════════════════════════════════════════════════════════════ +def stage2_crop_drawing(bgr: np.ndarray) -> Tuple[np.ndarray, str]: + """Remove right title block + bottom info strip via the dominant Canny line in + the outer regions. Never crops more than 40% of either axis.""" + h, w = bgr.shape[:2] + gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY) + edges = cv2.Canny(gray, 50, 150) + + right_zone = edges[:, int(w * 0.65):] + col_density = np.sum(right_zone > 0, axis=0) + crop_x = w + if col_density.max() > h * 0.4: + cand = np.where(col_density > h * 0.4)[0] + if len(cand) > 0: + crop_x = int(w * 0.65) + int(cand[0]) - 5 + + bottom_zone = edges[int(h * 0.75):, :] + row_density = np.sum(bottom_zone > 0, axis=1) + crop_y = h + if row_density.max() > w * 0.4: + cand = np.where(row_density > w * 0.4)[0] + if len(cand) > 0: + crop_y = int(h * 0.75) + int(cand[0]) - 5 + + crop_x = max(crop_x, int(w * 0.6)) + crop_y = max(crop_y, int(h * 0.6)) + cropped = bgr[:crop_y, :crop_x].copy() + return cropped, (f"[Stage 2] Cropped {cropped.shape[1]}x{cropped.shape[0]} " + f"(removed {w - crop_x}px right, {h - crop_y}px bottom)") + + +def _count_color_layers(bgr: np.ndarray, erase_mask: np.ndarray) -> int: + if not np.any(erase_mask): + return 0 + hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV) + hues = hsv[:, :, 0][erase_mask] + if len(hues) == 0: + return 0 + hist, _ = np.histogram(hues, bins=6, range=(0, 180)) + return int(np.sum(hist > len(hues) * 0.05)) + + +def _detect_colored_architecture(bgr: np.ndarray) -> Tuple[bool, float]: + """True when the WALLS are a muted colored (tan/orange) layer while MEP is + saturated — signalled by a large fraction of muted-chromatic non-paper pixels.""" + hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV) + S, V = hsv[:, :, 1], hsv[:, :, 2] + near_white = (S <= NEAR_WHITE_S_MAX) & (V >= NEAR_WHITE_V_MIN) + non_bg = int(np.count_nonzero(~near_white)) + if non_bg == 0: + return False, 0.0 + muted = ((S >= SAT_WALL_MIN) & (S < SAT_MEP_MIN) + & (V >= WALL_COLOR_V_MIN) & (V <= WALL_COLOR_V_MAX)) + ratio = float(np.count_nonzero(muted)) / non_bg + return ratio >= COLORED_ARCH_RATIO_MIN, ratio + + +def _chromatic_wall_mask(bgr: np.ndarray) -> np.ndarray: + hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV) + S, V = hsv[:, :, 1], hsv[:, :, 2] + mask = (S >= SAT_WALL_MIN) & (V >= WALL_COLOR_V_MIN) & (V <= WALL_COLOR_V_MAX) + return mask.astype(np.uint8) * 255 + + +def stage3_strip_colors(bgr: np.ndarray, chroma_threshold: int = 25, + strip_mode: str = "all") -> Tuple[np.ndarray, str]: + """Remove MEP color overlays. 'all' = erase every chromatic pixel (walls black) + or grayscale a monochromatic print. 'mep' = erase only saturated MEP, keep the + muted colored wall layer (colored-architecture sheets).""" + if strip_mode == "mep": + hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV) + S, V = hsv[:, :, 1], hsv[:, :, 2] + mep = ((S >= SAT_MEP_MIN) & (V > WALL_COLOR_V_MIN)).astype(np.uint8) * 255 + mep = cv2.dilate(mep, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)), iterations=1) + result = bgr.copy() + result[mep > 0] = (255, 255, 255) + return result, (f"[Stage 3] MEP-only strip: {int(np.count_nonzero(mep)):,} px, " + f"kept muted wall layer") + + img = bgr.astype(np.int32) + b, g, r = img[:, :, 0], img[:, :, 1], img[:, :, 2] + chroma = np.maximum(np.maximum(r, g), b) - np.minimum(np.minimum(r, g), b) + gray_arr = (0.299 * r + 0.587 * g + 0.114 * b).astype(np.int32) + colored_mask = (chroma > chroma_threshold) & (gray_arr < 240) + colored_pct = float(np.mean(colored_mask)) * 100 + + monochrome, dominant_hue = False, None + if colored_pct > 5.0: + hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV) + hues = hsv[:, :, 0][colored_mask] + if len(hues) > 0: + hist, edges = np.histogram(hues, bins=12, range=(0, 180)) + top = int(np.argmax(hist)) + if float(hist[top]) / float(np.sum(hist)) > 0.55: + monochrome = True + dominant_hue = int((edges[top] + edges[top + 1]) / 2) + + if monochrome: + result = cv2.cvtColor(cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY), cv2.COLOR_GRAY2BGR) + return result, f"[Stage 3] Monochromatic (hue~{dominant_hue}), kept grayscale" + result = bgr.copy() + result[colored_mask] = (255, 255, 255) + layers = _count_color_layers(bgr, colored_mask) + return result, (f"[Stage 3] Removed {layers} MEP color layers " + f"({np.count_nonzero(colored_mask):,} px erased)") + + +def _geometric_noise_filter(wall_mask: np.ndarray, min_area: int = 250, + min_extent: int = 40) -> int: + """Kill CCs where BOTH area < min_area AND max extent < min_extent. In-place.""" + n_lbl, labels, stats, _ = cv2.connectedComponentsWithStats(wall_mask, connectivity=8) + if n_lbl <= 1: + return 0 + killed = 0 + keep = np.ones(n_lbl, dtype=np.uint8) * 255 + keep[0] = 0 + for i in range(1, n_lbl): + area = int(stats[i, cv2.CC_STAT_AREA]) + bw, bh = int(stats[i, cv2.CC_STAT_WIDTH]), int(stats[i, cv2.CC_STAT_HEIGHT]) + if area < min_area and max(bw, bh) < min_extent: + keep[i] = 0 + killed += 1 + wall_mask[:] = keep[labels] + return killed + + +def _looks_like_closed_circle(wall_mask: np.ndarray, cx: int, cy: int, r: int) -> bool: + h, w = wall_mask.shape + a = np.linspace(0, 2 * np.pi, 64, endpoint=False) + xs = (cx + r * np.cos(a)).astype(int) + ys = (cy + r * np.sin(a)).astype(int) + valid = (xs >= 0) & (xs < w) & (ys >= 0) & (ys < h) + if np.sum(valid) < 48: + return False + return float(np.mean(wall_mask[ys[valid], xs[valid]] > 0)) >= 0.70 + + +def _erase_grid_bubbles(wall_mask: np.ndarray, bgr_source: np.ndarray) -> int: + """HoughCircles → erase grid-reference bubbles (circle + interior). In-place.""" + h, w = wall_mask.shape + blurred = cv2.GaussianBlur(cv2.cvtColor(bgr_source, cv2.COLOR_BGR2GRAY), (5, 5), 1.2) + min_r = max(8, int(min(h, w) * 0.004)) + max_r = max(min_r + 5, int(min(h, w) * 0.012)) + circles = cv2.HoughCircles(blurred, cv2.HOUGH_GRADIENT, dp=1.2, minDist=min_r * 3, + param1=80, param2=28, minRadius=min_r, maxRadius=max_r) + if circles is None: + return 0 + erased = 0 + for (cx, cy, r) in np.round(circles[0]).astype(int): + if not (r < cx < w - r and r < cy < h - r): + continue + if not _looks_like_closed_circle(wall_mask, cx, cy, r): + continue + cv2.circle(wall_mask, (cx, cy), r + 3, 0, -1) + erased += 1 + return erased + + +def _local_thickness(wall_mask: np.ndarray, point: Tuple[int, int]) -> int: + h, w = wall_mask.shape + x = max(0, min(w - 1, point[0])) + y = max(0, min(h - 1, point[1])) + if wall_mask[y, x] == 0: + return 3 + spans = [] + for dx, dy in [(1, 0), (0, 1), (1, 1), (1, -1)]: + span = 1 + for sign in (-1, 1): + for d in range(1, 30): + nx, ny = x + sign * d * dx, y + sign * d * dy + if not (0 <= nx < w and 0 <= ny < h) or wall_mask[ny, nx] == 0: + break + span += 1 + spans.append(span) + return int(np.median(spans)) + + +def _looks_like_door_arc(wall_mask: np.ndarray, cx: int, cy: int, r: int) -> bool: + h, w = wall_mask.shape + if not (r < cx < w - r and r < cy < h - r): + return False + a = np.linspace(0, 2 * np.pi, 64, endpoint=False) + xs = (cx + r * np.cos(a)).astype(int) + ys = (cy + r * np.sin(a)).astype(int) + valid = (xs >= 0) & (xs < w) & (ys >= 0) & (ys < h) + if np.sum(valid) < 32: + return False + hits = wall_mask[ys[valid], xs[valid]] > 0 + if not (0.10 <= float(np.mean(hits)) <= 0.55): + return False + presence = hits.astype(np.int8) + if len(presence) == 0: + return False + longest = run = 0 + for v in np.concatenate([presence, presence]): + run = run + 1 if v else 0 + longest = max(longest, run) + longest = min(longest, len(presence)) + return longest / len(presence) >= 0.12 + + +def _find_arc_endpoints(wall_mask: np.ndarray, cx: int, cy: int, + r: int) -> Optional[Tuple[Tuple[int, int], Tuple[int, int]]]: + h, w = wall_mask.shape + a = np.linspace(0, 2 * np.pi, 360, endpoint=False) + xs = np.clip((cx + r * np.cos(a)).astype(int), 0, w - 1) + ys = np.clip((cy + r * np.sin(a)).astype(int), 0, h - 1) + presence = wall_mask[ys, xs] > 0 + trans = np.diff(presence.astype(np.int8)) + starts, ends = np.where(trans == 1)[0], np.where(trans == -1)[0] + if len(starts) == 0 or len(ends) == 0: + return None + runs = [] + for s in starts: + ec = ends[ends > s] + if len(ec) > 0: + runs.append((s, ec[0], ec[0] - s)) + if not runs: + return None + runs.sort(key=lambda x: -x[2]) + bs, be, _ = runs[0] + return (int(xs[bs]), int(ys[bs])), (int(xs[be]), int(ys[be])) + + +def _close_door_arcs(wall_mask: np.ndarray, bgr_source: np.ndarray) -> int: + """Detect quarter-arc door swings, draw a chord across each opening. In-place.""" + blurred = cv2.GaussianBlur(cv2.cvtColor(bgr_source, cv2.COLOR_BGR2GRAY), (5, 5), 1.5) + h, w = wall_mask.shape + min_r = max(8, int(min(h, w) * 0.005)) + max_r = max(min_r + 5, int(min(h, w) * 0.05)) + circles = cv2.HoughCircles(blurred, cv2.HOUGH_GRADIENT, dp=1.2, minDist=min_r, + param1=70, param2=22, minRadius=min_r, maxRadius=max_r) + if circles is None: + return 0 + closed = 0 + for (cx, cy, r) in np.round(circles[0]).astype(int): + if not _looks_like_door_arc(wall_mask, cx, cy, r): + continue + eps = _find_arc_endpoints(wall_mask, cx, cy, r) + if eps is None: + continue + p1, p2 = eps + stroke = max(5, int(np.median([_local_thickness(wall_mask, p1), + _local_thickness(wall_mask, p2)]))) + cv2.line(wall_mask, p1, p2, 255, stroke) + closed += 1 + return closed + + +# ───────────────────────────────────────────────────────────────────────────── +# Step 10 — VECTOR WALL REFINEMENT (inlined port of vector_walls.refine_walls). +# Turns the heuristic pixel mask into structural walls: Hough segments → Manhattan +# snap (drops off-axis text/leader stubs) → collinear merge → corner completion → +# re-rasterize at distance-transform-measured stroke. This is what makes the mask +# match DeepPlan exactly. Pure cv2/numpy. Kill-switch WALL_VECTOR_REFINE=0. +# ───────────────────────────────────────────────────────────────────────────── +_Segment = Tuple[float, float, float, float] + + +def _px_per_metre(dpi: int, scale_denominator: int) -> float: + dpi = max(1, int(dpi)) + d = max(1, int(scale_denominator)) + return max(1e-6, (dpi / 0.0254) / d) + + +def _vw_detect_segments(mask: np.ndarray, min_len_px: int, max_gap_px: int) -> List[_Segment]: + lines = cv2.HoughLinesP((mask > 0).astype(np.uint8) * 255, rho=1, theta=np.pi / 360.0, + threshold=max(20, min_len_px // 2), minLineLength=min_len_px, + maxLineGap=max_gap_px) + if lines is None: + return [] + return [tuple(map(float, l[0])) for l in lines] + + +def _vw_angle_deg(seg: _Segment) -> float: + x1, y1, x2, y2 = seg + return math.degrees(math.atan2(y2 - y1, x2 - x1)) % 180.0 + + +def _vw_snap(seg: _Segment, tol_deg: float, snap_diagonals: bool) -> Optional[_Segment]: + a = _vw_angle_deg(seg) + targets = [0.0, 90.0] + ([45.0, 135.0] if snap_diagonals else []) + best = min(targets, key=lambda t: min(abs(a - t), 180.0 - abs(a - t))) + if min(abs(a - best), 180.0 - abs(a - best)) > tol_deg: + return None + x1, y1, x2, y2 = seg + cx, cy = (x1 + x2) / 2.0, (y1 + y2) / 2.0 + half = math.hypot(x2 - x1, y2 - y1) / 2.0 + rad = math.radians(best) + dx, dy = math.cos(rad) * half, math.sin(rad) * half + return (cx - dx, cy - dy, cx + dx, cy + dy) + + +def _vw_line_frame(seg: _Segment) -> Tuple[float, float, float]: + x1, y1, x2, y2 = seg + theta = math.atan2(y2 - y1, x2 - x1) + return theta, math.cos(theta), math.sin(theta) + + +def _vw_merge_collinear(segments: List[_Segment], perp_tol_px: float, + bridge_gap_px: float) -> List[_Segment]: + buckets: Dict[Tuple[int, int], List[_Segment]] = {} + for seg in segments: + theta, ux, uy = _vw_line_frame(seg) + adeg = int(round(math.degrees(theta) % 180.0)) + x1, y1, _, _ = seg + perp = -uy * x1 + ux * y1 + key = (adeg, int(round(perp / max(1.0, perp_tol_px)))) + buckets.setdefault(key, []).append(seg) + merged: List[_Segment] = [] + for group in buckets.values(): + theta, ux, uy = _vw_line_frame(group[0]) + intervals = [] + for x1, y1, x2, y2 in group: + t1, t2 = ux * x1 + uy * y1, ux * x2 + uy * y2 + intervals.append((min(t1, t2), max(t1, t2))) + intervals.sort() + x0, y0, _, _ = group[0] + perp = -uy * x0 + ux * y0 + px, py = -uy * perp, ux * perp + cur_a, cur_b = intervals[0] + for a, b in intervals[1:]: + if a <= cur_b + bridge_gap_px: + cur_b = max(cur_b, b) + else: + merged.append((px + ux * cur_a, py + uy * cur_a, px + ux * cur_b, py + uy * cur_b)) + cur_a, cur_b = a, b + merged.append((px + ux * cur_a, py + uy * cur_a, px + ux * cur_b, py + uy * cur_b)) + return merged + + +def _vw_intersect(s1: _Segment, s2: _Segment) -> Optional[Tuple[float, float]]: + x1, y1, x2, y2 = s1 + x3, y3, x4, y4 = s2 + d = (x1 - x2) * (y3 - y4) - (y1 - y2) * (x3 - x4) + if abs(d) < 1e-6: + return None + px = ((x1 * y2 - y1 * x2) * (x3 - x4) - (x1 - x2) * (x3 * y4 - y3 * x4)) / d + py = ((x1 * y2 - y1 * x2) * (y3 - y4) - (y1 - y2) * (x3 * y4 - y3 * x4)) / d + return px, py + + +def _vw_complete_corners(segments: List[_Segment], corner_gap_px: float) -> List[_Segment]: + segs = [list(s) for s in segments] + n = len(segs) + for i in range(n): + ai = _vw_angle_deg(tuple(segs[i])) + for j in range(i + 1, n): + aj = _vw_angle_deg(tuple(segs[j])) + perp = abs(ai - aj) + perp = min(perp, 180.0 - perp) + if abs(perp - 90.0) > 25.0: + continue + ip = _vw_intersect(tuple(segs[i]), tuple(segs[j])) + if ip is None: + continue + ix, iy = ip + for s in (segs[i], segs[j]): + d0 = math.hypot(s[0] - ix, s[1] - iy) + d2 = math.hypot(s[2] - ix, s[3] - iy) + if min(d0, d2) > corner_gap_px: + continue + if d0 <= d2: + s[0], s[1] = ix, iy + else: + s[2], s[3] = ix, iy + return [tuple(s) for s in segs] + + +def _vw_measure_stroke(mask: np.ndarray, lo_px: int, hi_px: int) -> int: + dist = cv2.distanceTransform((mask > 0).astype(np.uint8), cv2.DIST_L2, 5) + vals = dist[dist > 0] + if vals.size == 0: + return max(1, lo_px) + core = vals[vals >= np.median(vals)] + stroke = int(round(2.0 * float(np.median(core)))) + return int(np.clip(stroke, lo_px, hi_px)) + + +def _vw_rasterize(segments: List[_Segment], shape: Tuple[int, int], stroke: int) -> np.ndarray: + out = np.zeros(shape, np.uint8) + for x1, y1, x2, y2 in segments: + cv2.line(out, (int(round(x1)), int(round(y1))), (int(round(x2)), int(round(y2))), + 255, thickness=max(1, stroke), lineType=cv2.LINE_8) + return out + + +def refine_walls(wall_mask: np.ndarray, dpi: int = 150, scale_denominator: int = 100, + angle_tol_deg: float = 8.0, snap_diagonals: bool = False, + keep_original_union: bool = False) -> Tuple[np.ndarray, Dict[str, object]]: + """Refine heuristic wall mask into a structural one. Physical thresholds: + min segment 0.15 m · bridge gap 0.90 m · parallel tol 0.05 m · corner reach + 0.30 m · thickness 0.05-0.60 m. Returns original mask unchanged on any error.""" + info: Dict[str, object] = {} + try: + h, w = wall_mask.shape[:2] + ppm = _px_per_metre(dpi, scale_denominator) + min_len_px = max(8, int(0.15 * ppm)) + bridge_gap_px = max(6, int(0.90 * ppm)) + perp_tol_px = max(2, int(0.05 * ppm)) + corner_gap_px = max(4, int(0.30 * ppm)) + lo_px = max(1, int(0.05 * ppm)) + hi_px = max(lo_px + 1, int(0.60 * ppm)) + + raw = _vw_detect_segments(wall_mask, min_len_px, max_gap_px=perp_tol_px * 3) + info["segments_detected"] = len(raw) + if not raw: + return wall_mask, info + snapped, dropped = [], 0 + for seg in raw: + s = _vw_snap(seg, angle_tol_deg, snap_diagonals) + if s is None: + dropped += 1 + else: + snapped.append(s) + info["segments_dropped_offaxis"] = dropped + if not snapped: + return wall_mask, info + merged = _vw_merge_collinear(snapped, perp_tol_px, bridge_gap_px) + info["segments_after_merge"] = len(merged) + completed = _vw_complete_corners(merged, corner_gap_px) + stroke = _vw_measure_stroke(wall_mask, lo_px, hi_px) + info["stroke_px"] = stroke + refined = _vw_rasterize(completed, (h, w), stroke) + if keep_original_union: + refined = cv2.bitwise_or(refined, (wall_mask > 0).astype(np.uint8) * 255) + return refined, info + except Exception as exc: + info["error"] = f"{type(exc).__name__}: {exc}" + return wall_mask, info + + +def stage4_extract_walls(color_stripped: np.ndarray, cropped: np.ndarray, + include_color_walls: bool = False, + dpi: int = 150, scale_denominator: int = 100, + refine: bool = False, run_ocr: bool = False) -> Tuple[np.ndarray, str]: + """Steps 1-9: binarize (CLAHE + adaptive) → distance-transform thin-line filter + → geometric noise filter → grid-bubble erase → door-arc closure → directional + MORPH_CLOSE → +3px thicken → in-room noise sweep.""" + notes: List[str] = [] + gray = cv2.cvtColor(color_stripped, cv2.COLOR_BGR2GRAY) + + # 1. Binarize + enhanced = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8)).apply(gray) + block = max(11, int(min(color_stripped.shape[:2]) * 0.03) | 1) + binary = cv2.adaptiveThreshold(enhanced, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, + cv2.THRESH_BINARY_INV, block, 5) + if include_color_walls: + cw = _chromatic_wall_mask(color_stripped) + binary = cv2.bitwise_or(binary, cw) + notes.append(f"color-walls: {int(np.count_nonzero(cw)):,} px") + + # 2. Distance-transform thin-line filter. DT floor is a half-thickness, so 3.0 + # keeps only CCs with a ~6px-or-wider stroke (was 2.0 / ~4px, which let MEP + # runs and dimension lines through). Matches pipeline.py. + dist = cv2.distanceTransform(binary, cv2.DIST_L2, 5) + n_lbl, labels, stats, _ = cv2.connectedComponentsWithStats(binary, connectivity=8) + if n_lbl > 1: + thick_labels = np.unique(labels[(dist >= 3.0)]) + keep = np.zeros(n_lbl, dtype=np.uint8) + for lbl_id in thick_labels: + if lbl_id == 0: + continue + area = stats[lbl_id, cv2.CC_STAT_AREA] + wb, hb = stats[lbl_id, cv2.CC_STAT_WIDTH], stats[lbl_id, cv2.CC_STAT_HEIGHT] + if area >= 50 and max(wb, hb) >= 20: + keep[lbl_id] = 255 + wall_mask = keep[labels].astype(np.uint8) + else: + wall_mask = binary + + # 3. OCR text-erase — paint text bboxes to 0 (matches segmentation-pipeline mask). + if run_ocr and _HAS_EASYOCR: + erased = _ocr_erase_text_on_wall(wall_mask, run_ocr_raw(color_stripped)) + notes.append(f"text-erased: {erased}") + else: + notes.append(f"text-erased: 0{'' if run_ocr else ' (OCR off)'}") + # 4. Geometric noise filter + notes.append(f"geom-noise-killed: {_geometric_noise_filter(wall_mask)}") + # 5. Grid bubbles + notes.append(f"grid-bubbles: {_erase_grid_bubbles(wall_mask, cropped)}") + # 6. Door arcs + notes.append(f"doors-closed: {_close_door_arcs(wall_mask, cropped)}") + # 7. Directional MORPH_CLOSE H=30 V=15 + wall_mask = cv2.morphologyEx(wall_mask, cv2.MORPH_CLOSE, + cv2.getStructuringElement(cv2.MORPH_RECT, (30, 1))) + wall_mask = cv2.morphologyEx(wall_mask, cv2.MORPH_CLOSE, + cv2.getStructuringElement(cv2.MORPH_RECT, (1, 15))) + # 8. Light thicken +3px H+V + wall_mask = cv2.dilate(wall_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (3, 1))) + wall_mask = cv2.dilate(wall_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (1, 3))) + # 9. In-room noise sweep (erode → drop small CCs → dilate) + eroded = cv2.erode(wall_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))) + n2, lbl2, st2, _ = cv2.connectedComponentsWithStats(eroded, connectivity=8) + if n2 > 1: + keep2 = np.ones(n2, dtype=np.uint8) * 255 + keep2[0] = 0 + killed = 0 + for i in range(1, n2): + area = int(st2[i, cv2.CC_STAT_AREA]) + bw, bh = int(st2[i, cv2.CC_STAT_WIDTH]), int(st2[i, cv2.CC_STAT_HEIGHT]) + if area < 300 and max(bw, bh) < 50: + keep2[i] = 0 + killed += 1 + wall_mask = cv2.dilate(keep2[lbl2].astype(np.uint8), + cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))) + notes.append(f"in-room-noise: {killed}") + + # 10. Vector refinement — the structural cleanup that matches DeepPlan exactly. + if refine and os.environ.get("WALL_VECTOR_REFINE", "1") == "1": + refined, vinfo = refine_walls(wall_mask, dpi=dpi, scale_denominator=scale_denominator) + wall_mask = refined + notes.append( + f"vector-refined (seg={vinfo.get('segments_detected', 0)}" + f"->{vinfo.get('segments_after_merge', 0)} stroke={vinfo.get('stroke_px', '?')}px)" + if "error" not in vinfo else f"vector-refine skipped ({vinfo['error']})" + ) + + px = int(np.count_nonzero(wall_mask)) + return wall_mask, (f"[Stage 4] Walls: {px:,} px ({100.0 * px / wall_mask.size:.1f}%) " + + " ".join(notes)) + + +def extract_walls_and_crop(bgr: np.ndarray, dpi: int = 150, scale_denominator: int = 100, + run_ocr: bool = False, refine: bool = False + ) -> Tuple[np.ndarray, np.ndarray, List[str]]: + """DeepPlan heuristic Stages 2-4: crop → color-strip → wall extraction. + Returns (cropped_bgr, wall_mask, log). + + `refine` (step 10, vector refinement) defaults OFF: it re-rasterizes every + segment at one measured stroke, which turns welded MEP linework into uniform + fat ribbons. Off keeps the heuristic mask's true stroke widths.""" + log: List[str] = [] + cropped, m2 = stage2_crop_drawing(bgr) + log.append(m2) + colored_arch, ca_ratio = _detect_colored_architecture(cropped) + log.append(f"[Stage 2b] colored-architecture: {colored_arch} (ratio={ca_ratio:.2f})") + color_stripped, m3 = stage3_strip_colors(cropped, strip_mode="mep" if colored_arch else "all") + log.append(m3) + wall_mask, m4 = stage4_extract_walls(color_stripped, cropped, include_color_walls=colored_arch, + dpi=dpi, scale_denominator=scale_denominator, run_ocr=run_ocr, + refine=refine) + log.append(m4) + return cropped, wall_mask, log + + +# ───────────────────────────────────────────────────────────────────────────── +# Inlined mask → polygon +# ───────────────────────────────────────────────────────────────────────────── +def mask_to_polygon(mask_uint8: np.ndarray, min_area: int = 50, + epsilon_factor: float = 0.01) -> List[List[int]]: + """Binary mask → list of simplified polygons, each a flat [x1,y1,x2,y2,...].""" + if mask_uint8 is None or mask_uint8.size == 0: + return [] + m = (mask_uint8 > 0).astype(np.uint8) + cnts, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + out: List[List[int]] = [] + for c in cnts: + if cv2.contourArea(c) < min_area: + continue + eps = epsilon_factor * cv2.arcLength(c, True) + approx = cv2.approxPolyDP(c, eps, True) + if len(approx) < 3: + continue + out.append([int(v) for pt in approx.reshape(-1, 2) for v in pt]) + return out + + +# ───────────────────────────────────────────────────────────────────────────── +# Helpers +# ───────────────────────────────────────────────────────────────────────────── +def _downscale_pil(pil: Image.Image, maxdim: int) -> Image.Image: + w, h = pil.size + long_side = max(w, h) + if maxdim and long_side > maxdim: + s = maxdim / float(long_side) + pil = pil.resize((max(1, int(w * s)), max(1, int(h * s))), Image.LANCZOS) + return pil + + +def _wall_overlay_amber(cropped_bgr: np.ndarray, wall_mask: np.ndarray) -> np.ndarray: + """Extracted walls rendered exactly like the frontend 'Extract walls' overlay: + wall pixels amber (255,140,0) at ~59% opacity over the blueprint, native size. + Matches useWallMaskOverlay.js WALL_RGBA = [255,140,0,150].""" + out = cropped_bgr.astype(np.float32) + amber_bgr = np.array((0, 140, 255), dtype=np.float32) # RGB(255,140,0)→BGR + a = 150.0 / 255.0 + m = wall_mask > 0 + out[m] = (1 - a) * out[m] + a * amber_bgr + return cv2.cvtColor(out.astype(np.uint8), cv2.COLOR_BGR2RGB) + + +def _thick_walls_only(wall_mask: np.ndarray, min_stroke_px: int, + pct: float = 75.0) -> Tuple[np.ndarray, int, List[int]]: + """Keep only connected components whose wall stroke is >= min_stroke_px. + + SamAutomaticMaskGenerator cannot be prompted — it grid-samples whatever image + it is handed. The only way to tell it "these are the room boundaries" is to + remove everything else from the composite it sees. This drops thin linework + (partitions, MEP runs, dimension lines) so rooms are bounded by structural + walls alone; the cost is that rooms separated only by a thin wall merge. + + Stroke per CC = 2 x the `pct` percentile of the distance transform inside it. + A percentile rather than the max because wall junctions inflate the DT locally. + + Returns (filtered_mask, n_dropped, kept_strokes). min_stroke_px <= 0 is a no-op. + """ + if min_stroke_px <= 0: + return wall_mask, 0, [] + m = (wall_mask > 0).astype(np.uint8) + n_lbl, labels = cv2.connectedComponents(m, connectivity=8) + if n_lbl <= 1: + return wall_mask, 0, [] + dist = cv2.distanceTransform(m, cv2.DIST_L2, 5) + + sel = labels > 0 + lab_flat = labels[sel].ravel() + dist_flat = dist[sel].ravel() + order = np.argsort(lab_flat, kind="stable") + lab_flat, dist_flat = lab_flat[order], dist_flat[order] + ids = np.arange(1, n_lbl) + starts = np.searchsorted(lab_flat, ids, side="left") + ends = np.searchsorted(lab_flat, ids, side="right") + + keep = np.zeros(n_lbl, dtype=np.uint8) + dropped = 0 + strokes: List[int] = [] + for i, (s, e) in enumerate(zip(starts, ends), start=1): + if e <= s: + continue + stroke = 2.0 * float(np.percentile(dist_flat[s:e], pct)) + if stroke >= min_stroke_px: + keep[i] = 255 + strokes.append(int(round(stroke))) + else: + dropped += 1 + return keep[labels].astype(np.uint8), dropped, strokes + + +def _wall_composite(cropped_bgr: np.ndarray, wall_mask: np.ndarray, thicken: int) -> np.ndarray: + comp = cropped_bgr.copy() + wall = (wall_mask > 0).astype(np.uint8) + if thicken > 0: + k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * thicken + 1, 2 * thicken + 1)) + wall = cv2.dilate(wall, k, iterations=1) + comp[wall > 0] = (0, 0, 0) # BGR black walls + return comp + + +def _px_to_m2(area_px: float, scale_denom: int, dpi: int) -> float: + if not dpi or not scale_denom: + return 0.0 + m_per_px = (float(scale_denom) / float(dpi)) * 0.0254 + return area_px * (m_per_px ** 2) + + +def _iou(a_bool: np.ndarray, b_bool: np.ndarray) -> float: + inter = np.count_nonzero(a_bool & b_bool) + if inter == 0: + return 0.0 + return inter / float(np.count_nonzero(a_bool | b_bool)) + + +def _border_frac(mask_bool: np.ndarray, band: int = 3) -> float: + h, w = mask_bool.shape + edge = np.zeros((h, w), dtype=bool) + edge[:band, :] = edge[-band:, :] = edge[:, :band] = edge[:, -band:] = True + area = np.count_nonzero(mask_bool) + return np.count_nonzero(mask_bool & edge) / float(area) if area else 0.0 + + +# ═════════════════════════════════════════════════════════════════════════════ +# ENCLOSED-REGION DETECTION — the hard constraint SAM is not allowed to cross. +# +# A room is free space completely ringed by thick structural wall. Everything +# else (corridors bleeding off-sheet, shafts, half-open zones, drafting clutter) +# fails the enclosure test before SAM is ever asked about it. Free space is +# labelled 4-connected against 8-connected walls, so a mask cannot squeeze +# diagonally through a wall corner. +# ═════════════════════════════════════════════════════════════════════════════ +def _enclosed_regions(wall_thick: np.ndarray, seal_px: int, min_area_px: int, + max_area_px: int, enclosure_min: float + ) -> Tuple[np.ndarray, np.ndarray, int, List[Dict[str, Any]], Dict[str, int]]: + """Free-space components that are fully enclosed by thick wall. + + `seal_px` MORPH_CLOSEs the wall first, bridging door openings and small wall + gaps so a doorway does not fuse two rooms into one region. + `enclosure_min` is the fraction of a region's 1px outer ring that must be + wall. 1.0 demands a perfectly continuous boundary; ~0.98 tolerates a few + stray pixels without admitting a room with a genuine hole in its perimeter. + + Returns (wall_bool_sealed, free_labels, n_labels, regions, reject_counts). + """ + wall = (wall_thick > 0).astype(np.uint8) + if seal_px > 0: + k = cv2.getStructuringElement(cv2.MORPH_RECT, (2 * seal_px + 1, 2 * seal_px + 1)) + wall = cv2.morphologyEx(wall, cv2.MORPH_CLOSE, k) + wall_bool = wall > 0 + + free = (~wall_bool).astype(np.uint8) + n_lbl, labels, stats, cents = cv2.connectedComponentsWithStats(free, connectivity=4) + h, w = wall.shape + k3 = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3)) + + regions: List[Dict[str, Any]] = [] + rejects: Dict[str, int] = {} + + def _rej(why: str) -> None: + rejects[why] = rejects.get(why, 0) + 1 + + for i in range(1, n_lbl): + x, y, bw, bh, area = (int(v) for v in stats[i]) + # Touching the sheet edge means the region escapes the drawing — it is the + # exterior, or an open corridor running off-sheet. Never a closed room. + if x <= 0 or y <= 0 or x + bw >= w or y + bh >= h: + _rej("open / touches sheet edge") + continue + if area < min_area_px: + _rej("below min area") + continue + if area > max_area_px: + _rej("above max area") + continue + + y0, y1 = max(0, y - 2), min(h, y + bh + 2) + x0, x1 = max(0, x - 2), min(w, x + bw + 2) + sub = labels[y0:y1, x0:x1] == i + ring = cv2.dilate(sub.astype(np.uint8), k3, iterations=1).astype(bool) & ~sub + n_ring = int(np.count_nonzero(ring)) + if n_ring == 0: + _rej("degenerate") + continue + wall_frac = float(np.count_nonzero(ring & wall_bool[y0:y1, x0:x1])) / n_ring + if wall_frac < enclosure_min: + _rej("perimeter not continuous wall") + continue + + regions.append({ + "label": i, + "bbox": (x, y, bw, bh), + "area": area, + "slice": (y0, y1, x0, x1), + "sub": sub, + "centroid": (float(cents[i][0]), float(cents[i][1])), + "wall_frac": wall_frac, + }) + return wall_bool, labels, n_lbl, regions, rejects + + +def _label_deep_points(labels: np.ndarray, n_lbl: int, + dist: np.ndarray) -> Dict[int, Tuple[int, int]]: + """Deepest interior pixel (max distance-to-wall) per free-space label.""" + sel = labels > 0 + if not np.any(sel): + return {} + ys, xs = np.nonzero(sel) + lab_flat = labels[sel].ravel() + dist_flat = dist[sel].ravel() + order = np.argsort(lab_flat, kind="stable") + lab_flat, dist_flat = lab_flat[order], dist_flat[order] + ys, xs = ys[order], xs[order] + ids = np.arange(1, n_lbl) + starts = np.searchsorted(lab_flat, ids, side="left") + ends = np.searchsorted(lab_flat, ids, side="right") + out: Dict[int, Tuple[int, int]] = {} + for i, (s, e) in enumerate(zip(starts, ends), start=1): + if e <= s: + continue + j = s + int(np.argmax(dist_flat[s:e])) + out[i] = (int(xs[j]), int(ys[j])) + return out + + +def _positive_points(sub: np.ndarray, k: int) -> List[Tuple[int, int]]: + """Up to k positive seeds at successive distance-transform maxima. + + The first is the deepest point in the region — maximally far from every wall, + which is what keeps a seed off a boundary in an L- or T-shaped room. Each + pick suppresses a disc of its own radius so the next lands in a different + limb of the shape rather than beside the first. + """ + d = cv2.distanceTransform(sub.astype(np.uint8), cv2.DIST_L2, 5) + work = d.copy() + pts: List[Tuple[int, int]] = [] + for _ in range(max(1, k)): + _, mx, _, loc = cv2.minMaxLoc(work) + if mx <= 0: + break + pts.append((int(loc[0]), int(loc[1]))) + cv2.circle(work, (int(loc[0]), int(loc[1])), max(3, int(mx)), 0, -1) + return pts + + +def _negative_points(region: Dict[str, Any], labels: np.ndarray, wall_bool: np.ndarray, + deep_by_label: Dict[int, Tuple[int, int]], + n_wall: int, n_neigh: int, reach: int) -> List[Tuple[int, int]]: + """Negatives that fence the region in. + + Two kinds, both aimed at leakage: + (a) on the wall ring itself — corridors, shafts, door jambs and wall gaps + all present as wall pixels bordering the room, so this marks the + boundary as not-room from every side. + (b) the deep point of each adjacent free region — a doorway or wall gap is + exactly where SAM would spill into the neighbour, and a negative sitting + in the middle of that neighbour is the cheapest way to say "not there". + The exterior is one of these regions, so this also fences the outside. + """ + y0, y1, x0, x1 = region["slice"] + r = max(3, int(reach)) + # Window padded by the full reach. `region["slice"]` is only bbox+2px, and + # dilating inside it clips the ring to 2px however large the reach — which + # caps it silently, so the ring never crosses a wall to touch the + # neighbouring room and the adjacent-region negatives never fire. + h, w = labels.shape[:2] + wy0, wy1 = max(0, y0 - r), min(h, y1 + r) + wx0, wx1 = max(0, x0 - r), min(w, x1 + r) + win = labels[wy0:wy1, wx0:wx1] + sub = win == region["label"] + + # 2r+1, not r: a rect kernel of side r dilates by r//2 in each direction, so + # sizing it r would deliver half the requested reach. + kr = cv2.getStructuringElement(cv2.MORPH_RECT, (2 * r + 1, 2 * r + 1)) + ring = cv2.dilate(sub.astype(np.uint8), kr, iterations=1).astype(bool) & ~sub + + negs: List[Tuple[int, int]] = [] + ys, xs = np.nonzero(ring & wall_bool[wy0:wy1, wx0:wx1]) + if xs.size and n_wall > 0: + idx = np.linspace(0, xs.size - 1, min(n_wall, xs.size)).astype(int) + negs += [(int(xs[j]) + wx0, int(ys[j]) + wy0) for j in idx] + + if n_neigh > 0: + neigh = [int(v) for v in np.unique(win[ring]) if v > 0 and v != region["label"]] + for nb in neigh[:n_neigh]: + p = deep_by_label.get(nb) + if p is not None: + negs.append(p) + return negs + + +# ───────────────────────────────────────────────────────────────────────────── +# Geometric validation — only simple closed room polygons survive +# ───────────────────────────────────────────────────────────────────────────── +def _geom_valid(mask_bool: np.ndarray, min_solidity: float, min_extent: float, + max_vertices: int, min_axis_frac: float, frag_max: float, + eps_frac: float = 0.01, axis_tol_deg: float = 12.0 + ) -> Tuple[bool, str, Dict[str, float]]: + """Accept square / rectangle / L / T / other simple closed rectilinear rooms. + + solidity area / convex-hull area. A rectangle is 1.0, an L or T about 0.7, + a leaking or ragged mask much lower. + extent area / bbox area. Same idea, catches slivers and diagonals. + vertices after approxPolyDP. A room is 4-12 corners; a noisy blob is dozens. + axis_frac length-weighted fraction of the outline running within + axis_tol_deg of horizontal or vertical. Rooms are rectilinear; + masks that followed pipe runs or arcs are not. + frag area outside the largest contour. Non-zero means the mask is in + pieces, which is not one room. + """ + metrics: Dict[str, float] = {} + m = mask_bool.astype(np.uint8) + cnts, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + if not cnts: + return False, "empty mask", metrics + cnts = sorted(cnts, key=cv2.contourArea, reverse=True) + main = cnts[0] + a_main = float(cv2.contourArea(main)) + if a_main <= 0: + return False, "degenerate contour", metrics + + frag = float(sum(cv2.contourArea(c) for c in cnts[1:])) / a_main + hull_a = float(cv2.contourArea(cv2.convexHull(main))) + solidity = a_main / hull_a if hull_a > 0 else 0.0 + _, _, bw, bh = cv2.boundingRect(main) + extent = a_main / float(max(1, bw * bh)) + + peri = cv2.arcLength(main, True) + approx = cv2.approxPolyDP(main, eps_frac * peri, True).reshape(-1, 2) + n_vert = len(approx) + + total_len = axis_len = 0.0 + for j in range(n_vert): + p, q = approx[j], approx[(j + 1) % n_vert] + dx, dy = float(q[0] - p[0]), float(q[1] - p[1]) + seg_len = math.hypot(dx, dy) + if seg_len <= 0: + continue + ang = math.degrees(math.atan2(dy, dx)) % 90.0 + total_len += seg_len + if min(ang, 90.0 - ang) <= axis_tol_deg: + axis_len += seg_len + axis_frac = axis_len / total_len if total_len > 0 else 0.0 + + metrics = {"solidity": round(solidity, 3), "extent": round(extent, 3), + "vertices": float(n_vert), "axis_frac": round(axis_frac, 3), + "frag": round(frag, 3)} + + if frag > frag_max: + return False, f"fragmented ({frag:.0%} of area outside main part)", metrics + if n_vert < 4: + return False, f"{n_vert} vertices (not a closed polygon)", metrics + if n_vert > max_vertices: + return False, f"{n_vert} vertices > {max_vertices} (noisy outline)", metrics + if solidity < min_solidity: + return False, f"solidity {solidity:.2f} < {min_solidity:.2f}", metrics + if extent < min_extent: + return False, f"extent {extent:.2f} < {min_extent:.2f}", metrics + if axis_frac < min_axis_frac: + return False, f"only {axis_frac:.0%} of outline axis-aligned", metrics + return True, "ok", metrics + + +# ───────────────────────────────────────────────────────────────────────────── +# Prompted room segmentation — one SAM call per enclosed region +# ───────────────────────────────────────────────────────────────────────────── +def _prompt_debug_image(base_rgb: np.ndarray, regions: List[Dict[str, Any]], + prompts: List[Dict[str, Any]]) -> np.ndarray: + """Green = positive point, red = negative, yellow box = region bbox prompt.""" + viz = base_rgb.copy() + for r in regions: + x, y, bw, bh = r["bbox"] + cv2.rectangle(viz, (x, y), (x + bw, y + bh), (255, 200, 0), 1) + rad = max(2, int(min(viz.shape[:2]) / 400)) + for p in prompts: + for (px, py) in p["neg"]: + cv2.circle(viz, (px, py), rad, (255, 40, 40), -1) + for (px, py) in p["pos"]: + cv2.circle(viz, (px, py), rad + 1, (0, 220, 60), -1) + return viz + + +def segment_rooms_prompted(predictor: Any, sam_rgb: np.ndarray, wall_thick: np.ndarray, + p: Dict[str, Any], log: List[str] + ) -> Tuple[List[Dict[str, Any]], np.ndarray]: + """Enclosed regions → point + box prompts → SAM → validate → merge. + + Replaces SamAutomaticMaskGenerator entirely. One image embedding for the whole + sheet, then a cheap mask-decoder call per candidate room, instead of ~341 + encoder passes over a recursive crop pyramid. + """ + h, w = wall_thick.shape[:2] + total = float(h * w) + max_area_px = int(p["max_area_frac"] * total) + + wall_bool, labels, n_lbl, regions, rejects = _enclosed_regions( + wall_thick, int(p["seal_px"]), int(p["min_area_px"]), max_area_px, + float(p["enclosure_min"])) + log.append(f"[rooms] enclosed regions: {len(regions)} accepted") + for why, cnt in sorted(rejects.items(), key=lambda kv: -kv[1]): + log.append(f"[rooms] rejected {cnt}: {why}") + if not regions: + return [], sam_rgb + + free_dist = cv2.distanceTransform((~wall_bool).astype(np.uint8), cv2.DIST_L2, 5) + deep_by_label = _label_deep_points(labels, n_lbl, free_dist) + + predictor.set_image(sam_rgb) + log.append("[rooms] image embedding computed once; decoding one mask per region") + + rooms: List[Dict[str, Any]] = [] + prompts: List[Dict[str, Any]] = [] + n_fallback = 0 + geom_rejects: Dict[str, int] = {} + + for reg in regions: + y0, y1, x0, x1 = reg["slice"] + pos_local = _positive_points(reg["sub"], int(p["n_pos"])) + if not pos_local: + continue + pos = [(px + x0, py + y0) for (px, py) in pos_local] + neg = _negative_points(reg, labels, wall_bool, deep_by_label, + int(p["n_neg_wall"]), int(p["n_neg_neigh"]), int(p["neg_reach"])) + + pts = np.array(pos + neg, dtype=np.float32) + lbls = np.array([1] * len(pos) + [0] * len(neg), dtype=np.int32) + bx, by, bw, bh = reg["bbox"] + pad = int(p["box_pad"]) + box = np.array([max(0, bx - pad), max(0, by - pad), + min(w, bx + bw + pad), min(h, by + bh + pad)], dtype=np.float32) + prompts.append({"pos": pos, "neg": neg}) + + try: + masks, scores, _ = predictor.predict( + point_coords=pts, point_labels=lbls, + box=box if p["use_box"] else None, + multimask_output=True, + ) + except Exception as exc: + log.append(f"[rooms] SAM predict failed on region {reg['label']}: " + f"{type(exc).__name__}: {exc}") + continue + + region_full = np.zeros((h, w), dtype=bool) + region_full[y0:y1, x0:x1] = reg["sub"] + other_free = (labels > 0) & (labels != reg["label"]) + sx, sy = pos[0] + + best = None + for cand, sam_score in zip(masks, scores): + mb = cand.astype(bool) & ~wall_bool # hard constraint: never cross wall + if p["hard_clip"]: + # Keep only the piece connected to the seed. After clipping at the + # wall, anything reachable from the seed is inside this region by + # construction, so leakage becomes structurally impossible rather + # than merely penalised. + n_cc, cc = cv2.connectedComponents(mb.astype(np.uint8), connectivity=4) + if n_cc <= 1: + continue + sid = int(cc[sy, sx]) + if sid == 0: + continue + mb = cc == sid + area = int(np.count_nonzero(mb)) + if area == 0: + continue + leak = float(np.count_nonzero(mb & other_free)) / area + iou = _iou(mb, region_full) + obj = iou - float(p["leak_penalty"]) * leak + if best is None or obj > best["obj"]: + best = {"mask": mb, "iou": iou, "leak": leak, "obj": obj, + "sam": float(sam_score), "area": area} + + used_fallback = False + if best is None or best["iou"] < float(p["min_region_iou"]) \ + or best["leak"] > float(p["max_leak"]): + if not p["fallback_region"]: + geom_rejects["SAM mask failed IoU/leak gate"] = \ + geom_rejects.get("SAM mask failed IoU/leak gate", 0) + 1 + continue + # The region is already proven enclosed, so it is a valid room even + # when SAM's own mask is not. Take the region and say so. + best = {"mask": region_full, "iou": 1.0, "leak": 0.0, "obj": 1.0, + "sam": 0.0, "area": reg["area"]} + used_fallback = True + n_fallback += 1 + + ok, why, metrics = _geom_valid( + best["mask"], float(p["min_solidity"]), float(p["min_extent"]), + int(p["max_vertices"]), float(p["min_axis_frac"]), float(p["frag_max"])) + if not ok: + geom_rejects[why.split(" (")[0]] = geom_rejects.get(why.split(" (")[0], 0) + 1 + continue + + ys, xs = np.nonzero(best["mask"]) + rooms.append({ + "segmentation": best["mask"], + "area": int(best["area"]), + "bbox": (int(xs.min()), int(ys.min()), + int(xs.max() - xs.min() + 1), int(ys.max() - ys.min() + 1)), + "score": best["obj"], + "iou_region": round(best["iou"], 3), + "leak": round(best["leak"], 4), + "sam_score": round(best["sam"], 3), + "fallback": used_fallback, + "metrics": metrics, + }) + + for why, cnt in sorted(geom_rejects.items(), key=lambda kv: -kv[1]): + log.append(f"[rooms] dropped {cnt}: {why}") + if n_fallback: + log.append(f"[rooms] {n_fallback} region(s) kept as-is — SAM mask missed the " + f"enclosure, the region itself is already wall-bounded") + + rooms = _merge_rooms(rooms, float(p["merge_iou"]), float(p["contain_max"])) + rooms.sort(key=lambda r: (r["bbox"][1], r["bbox"][0])) + + viz = _prompt_debug_image(sam_rgb, regions, prompts) + return rooms, viz + + +def _merge_rooms(rooms: List[Dict[str, Any]], merge_iou: float, + contain_max: float) -> List[Dict[str, Any]]: + """Fuse over-segmented masks of one room; drop under-segmented duplicates. + + Regions are disjoint by construction, so this mostly catches the case where + SAM returned near-identical masks for two seeds in the same space. Two masks + merge when they overlap by more than `merge_iou`; a mask that sits more than + `contain_max` inside an already-kept one is a duplicate and is dropped. + """ + if not rooms: + return [] + rooms = sorted(rooms, key=lambda r: r["area"], reverse=True) + kept: List[Dict[str, Any]] = [] + for r in rooms: + merged = False + for k in kept: + inter = int(np.count_nonzero(r["segmentation"] & k["segmentation"])) + if inter == 0: + continue + if inter / float(r["area"]) > contain_max: + merged = True # duplicate / subset — discard + break + if _iou(r["segmentation"], k["segmentation"]) > merge_iou: + k["segmentation"] = k["segmentation"] | r["segmentation"] + ys, xs = np.nonzero(k["segmentation"]) + k["area"] = int(np.count_nonzero(k["segmentation"])) + k["bbox"] = (int(xs.min()), int(ys.min()), + int(xs.max() - xs.min() + 1), int(ys.max() - ys.min() + 1)) + merged = True + break + if not merged: + kept.append(r) + return kept + + +# ───────────────────────────────────────────────────────────────────────────── +# Rendering +# ───────────────────────────────────────────────────────────────────────────── +def _color_overlay(cropped_bgr: np.ndarray, rooms: List[Dict[str, Any]]) -> np.ndarray: + seg = cropped_bgr.copy() + for i, r in enumerate(rooms): + rc = ROOM_COLORS[i % len(ROOM_COLORS)] + col = np.array((rc[2], rc[1], rc[0]), dtype=np.float32) + m = r["segmentation"] + seg[m] = (0.45 * col + 0.55 * seg[m]).astype(np.uint8) + for i, r in enumerate(rooms): + M = cv2.moments(r["segmentation"].astype(np.uint8)) + if M["m00"]: + cx, cy = int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]) + cv2.putText(seg, str(i + 1), (cx, cy), cv2.FONT_HERSHEY_SIMPLEX, + max(0.5, seg.shape[1] / 2200.0), (0, 0, 0), 2, cv2.LINE_AA) + return cv2.cvtColor(seg, cv2.COLOR_BGR2RGB) + + +def _boundary_overlay(cropped_bgr: np.ndarray, rooms: List[Dict[str, Any]]) -> np.ndarray: + seg = cropped_bgr.copy() + for i, r in enumerate(rooms): + rc = ROOM_COLORS[i % len(ROOM_COLORS)] + cnts, _ = cv2.findContours(r["segmentation"].astype(np.uint8), + cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + cv2.drawContours(seg, cnts, -1, (rc[2], rc[1], rc[0]), 2, cv2.LINE_AA) + return cv2.cvtColor(seg, cv2.COLOR_BGR2RGB) + + +def _instances(cropped_bgr: np.ndarray, rooms: List[Dict[str, Any]]) -> List[np.ndarray]: + out = [] + for r in rooms: + x, y, w, h = r["bbox"] + sub = cropped_bgr[y:y + h, x:x + w].copy() + m = r["segmentation"][y:y + h, x:x + w] + sub[~m] = (sub[~m] * 0.25).astype(np.uint8) + out.append(cv2.cvtColor(sub, cv2.COLOR_BGR2RGB)) + return out + + +# ───────────────────────────────────────────────────────────────────────────── +# Export +# ───────────────────────────────────────────────────────────────────────────── +def _room_records(rooms: List[Dict[str, Any]], scale_denom: int, dpi: int) -> List[Dict[str, Any]]: + recs = [] + for i, r in enumerate(rooms): + polys = mask_to_polygon(r["segmentation"].astype(np.uint8)) + M = cv2.moments(r["segmentation"].astype(np.uint8)) + cx = int(M["m10"] / M["m00"]) if M["m00"] else int(r["bbox"][0]) + cy = int(M["m01"] / M["m00"]) if M["m00"] else int(r["bbox"][1]) + recs.append({ + "id": i + 1, "area_px": int(r["area"]), + "area_m2": round(_px_to_m2(r["area"], scale_denom, dpi), 3), + "bbox": list(r["bbox"]), "centroid": [cx, cy], "polygons": polys, + }) + return recs + + +def _export_files(overlay_rgb: np.ndarray, recs: List[Dict[str, Any]], + size: Tuple[int, int]) -> Tuple[str, str, str]: + ts = int(time.time()) + png_path = os.path.join(_EXPORT_DIR, f"rooms_{ts}.png") + svg_path = os.path.join(_EXPORT_DIR, f"rooms_{ts}.svg") + json_path = os.path.join(_EXPORT_DIR, f"rooms_{ts}.json") + + Image.fromarray(overlay_rgb).save(png_path) + + w, h = size + parts = [f''] + for r in recs: + rc = ROOM_COLORS[(r["id"] - 1) % len(ROOM_COLORS)] + fill = f"rgb({rc[0]},{rc[1]},{rc[2]})" + for poly in r["polygons"]: + pts = " ".join(f"{poly[j]},{poly[j + 1]}" for j in range(0, len(poly) - 1, 2)) + parts.append(f'') + parts.append(f'{r["id"]}') + parts.append("") + with open(svg_path, "w", encoding="utf-8") as f: + f.write("\n".join(parts)) + + with open(json_path, "w", encoding="utf-8") as f: + json.dump({"rooms": recs, "count": len(recs)}, f, indent=2) + return png_path, svg_path, json_path + + +# ───────────────────────────────────────────────────────────────────────────── +# Gradio callback +# ───────────────────────────────────────────────────────────────────────────── +def _walls_from_pdf(pdf_file: Any, page_no: int, render_dpi: int, scale_denom: int, + min_thk_mm: float, max_thk_mm: float, min_wall_len_mm: float, + use_strokes: bool, use_pairs: bool, use_fills: bool, + log: List[str]) -> Optional[Tuple[np.ndarray, np.ndarray, int]]: + """Vector path: PDF page → (cropped BGR, wall mask, effective dpi). + + Returns None when it cannot apply, having said why in the log, so the caller + falls back to the raster pipeline instead of failing. + """ + if pdf_file is None: + return None + if not _HAS_PDF: + log.append(f"[pdf] vector extraction unavailable ({_PDF_ERR}) — using raster path") + return None + path = getattr(pdf_file, "name", None) or str(pdf_file) + try: + with open(path, "rb") as fh: + data = fh.read() + doc = pdf_walls.open_document(data) + except Exception as exc: + log.append(f"[pdf] {exc} — using raster path") + return None + + idx = max(0, min(int(page_no) - 1, doc.page_count - 1)) + if idx != int(page_no) - 1: + log.append(f"[pdf] page {page_no} out of range; using page {idx + 1} of {doc.page_count}") + page = doc[idx] + + try: + mask, info, plog = pdf_walls.extract_wall_mask( + page, dpi=int(render_dpi), scale_denominator=int(scale_denom), + min_thickness_mm=float(min_thk_mm), max_thickness_mm=float(max_thk_mm), + min_wall_length_mm=float(min_wall_len_mm), + use_strokes=bool(use_strokes), use_pairs=bool(use_pairs), + use_fills=bool(use_fills)) + log += plog + if not info.get("vector") or info.get("wall_px", 0) == 0: + log.append("[pdf] falling back to the raster pipeline") + return None + page_bgr = pdf_walls.render_page(page, int(render_dpi)) + except Exception as exc: + log.append(f"[pdf] vector extraction failed ({type(exc).__name__}: {exc}) — using raster path") + return None + + # Same title-block crop the raster path applies, derived from the rendered + # page so the mask and the image stay in one coordinate space. + cropped, m2 = stage2_crop_drawing(page_bgr) + log.append(m2) + mask = mask[:cropped.shape[0], :cropped.shape[1]] + if mask.shape[:2] != cropped.shape[:2]: + mask = cv2.resize(mask, (cropped.shape[1], cropped.shape[0]), + interpolation=cv2.INTER_NEAREST) + log.append("[pdf] source: VECTOR geometry — thresholds are physical (mm), " + "not pixels, so render DPI does not affect the result") + return cropped, mask, int(render_dpi) + + +def segment(image: Optional[Image.Image], + pdf_file: Any, pdf_page: int, pdf_dpi: int, + pdf_min_thk: float, pdf_max_thk: float, pdf_min_len: float, + pdf_strokes: bool, pdf_pairs: bool, pdf_fills: bool, + maxdim: int, scale_denom: int, dpi: int, + ocr_erase: bool, sam_min_stroke: int, seal_px: int, wall_thicken: int, + n_pos: int, n_neg_wall: int, n_neg_neigh: int, neg_reach: int, + box_pad: int, use_box: bool, hard_clip: bool, + min_area_px: int, max_area_frac: float, enclosure_min: float, + min_region_iou: float, max_leak: float, leak_penalty: float, + fallback_region: bool, + min_solidity: float, min_extent: float, max_vertices: int, + min_axis_frac: float, frag_max: float, + merge_iou: float, contain_max: float): + if image is None and pdf_file is None: + return (None, None, None, None, None, [], [], [], + "Upload a floor plan image, or a PDF for the vector path.", + None, None, None) + t0 = time.time() + log: List[str] = [] + + # Vector first when a PDF is supplied. A vector PDF carries the wall geometry + # exactly — coordinates, stroke widths, fills — so there is nothing to recover + # by thresholding pixels, and no render DPI for the thresholds to depend on. + vec = _walls_from_pdf(pdf_file, pdf_page, pdf_dpi, scale_denom, + pdf_min_thk, pdf_max_thk, pdf_min_len, + pdf_strokes, pdf_pairs, pdf_fills, log) + if vec is not None: + cropped, wall_mask, eff_dpi = vec + eff_scale = int(scale_denom) if scale_denom else 100 + return _segment_from_walls(cropped, wall_mask, eff_dpi, eff_scale, t0, log, + maxdim, sam_min_stroke, seal_px, wall_thicken, + n_pos, n_neg_wall, n_neg_neigh, neg_reach, + box_pad, use_box, hard_clip, + min_area_px, max_area_frac, enclosure_min, + min_region_iou, max_leak, leak_penalty, + fallback_region, min_solidity, min_extent, + max_vertices, min_axis_frac, frag_max, + merge_iou, contain_max) + if image is None: + return (None, None, None, None, None, [], [], [], + "\n".join(log + ["No usable vector geometry and no image uploaded."]), + None, None, None) + + # Wall extraction in this app is VECTOR — NOT the heuristic Stages 2-4. The + # blueprint is cropped to drop the title block, then traced to vector walls + # locally by vtracer; the downstream thick-stroke filter keeps structural walls + # only. Downscale only as a memory guard for oversized sheets. + src_pil = image + if EXTRACT_MAX_DIM and max(image.size) > EXTRACT_MAX_DIM: + src_pil = _downscale_pil(image.convert("RGB"), EXTRACT_MAX_DIM) + log.append(f"[walls] image > {EXTRACT_MAX_DIM}px — downscaled for memory") + bgr = cv2.cvtColor(np.array(src_pil.convert("RGB")), cv2.COLOR_RGB2BGR) + eff_scale = int(scale_denom) if scale_denom else 100 + # Wall extraction is defined in pixels and ignores this; DPI only converts the + # finished room areas to m². + eff_dpi = int(dpi) if dpi else 150 + cropped, m2 = stage2_crop_drawing(bgr) + log.append(m2) + wall_mask, wall_source = extract_walls_primary(cropped, log, eff_dpi, eff_scale) + if wall_mask is None: + return (cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB), None, None, None, None, + [], [], [], + "\n".join(log + ["Vector wall extraction failed — no walls extracted."]), + None, None, None) + log.append(f"[walls] source: {wall_source}") + return _segment_from_walls(cropped, wall_mask, eff_dpi, eff_scale, t0, log, + maxdim, sam_min_stroke, seal_px, wall_thicken, + n_pos, n_neg_wall, n_neg_neigh, neg_reach, + box_pad, use_box, hard_clip, + min_area_px, max_area_frac, enclosure_min, + min_region_iou, max_leak, leak_penalty, + fallback_region, min_solidity, min_extent, + max_vertices, min_axis_frac, frag_max, + merge_iou, contain_max) + + +def _segment_from_walls(cropped: np.ndarray, wall_mask: np.ndarray, + eff_dpi: int, eff_scale: int, t0: float, log: List[str], + maxdim: int, sam_min_stroke: int, seal_px: int, wall_thicken: int, + n_pos: int, n_neg_wall: int, n_neg_neigh: int, neg_reach: int, + box_pad: int, use_box: bool, hard_clip: bool, + min_area_px: int, max_area_frac: float, enclosure_min: float, + min_region_iou: float, max_leak: float, leak_penalty: float, + fallback_region: bool, + min_solidity: float, min_extent: float, max_vertices: int, + min_axis_frac: float, frag_max: float, + merge_iou: float, contain_max: float): + """Everything downstream of "we have a wall mask" — shared by the vector and + raster paths, so both are segmented by identical code.""" + walls_rgb = _wall_overlay_amber(cropped, wall_mask) # frontend-style walls preview (full res) + + # Thin-line rejection. stage4's distance-transform filter already drops strokes + # under ~6px; this second pass measures each surviving component and keeps only + # genuinely structural walls, so sprinkler runs, electrical, plumbing, + # dimensions, text, symbols, furniture and drafting lines are gone before any + # region or prompt is derived. The unfiltered mask stays for the amber preview. + thick_mask = wall_mask + if int(sam_min_stroke) > 0: + thick_mask, n_thin, strokes = _thick_walls_only(wall_mask, int(sam_min_stroke)) + kept_desc = (f"kept {len(strokes)} (stroke {min(strokes)}-{max(strokes)}px)" + if strokes else "kept 0 — threshold too high, no walls left") + log.append(f"[walls] thick-wall filter >={int(sam_min_stroke)}px: " + f"dropped {n_thin} thin components, {kept_desc}") + composite = _wall_composite(cropped, thick_mask, int(wall_thicken)) + composite_rgb = cv2.cvtColor(composite, cv2.COLOR_BGR2RGB) + + # Downscale ONLY for SAM (cost ~quadratic in pixels). Rooms come back in this + # small frame, then get upscaled to the full cropped frame below. The thick + # wall mask is resized with MAX-pooling, not nearest — nearest drops isolated + # wall pixels and punches holes that break the enclosure test. + H, W = cropped.shape[:2] + sf = 1.0 + if max(H, W) > int(maxdim): + sf = int(maxdim) / float(max(H, W)) + sw, sh = max(1, int(W * sf)), max(1, int(H * sf)) + sam_input = cv2.resize(composite_rgb, (sw, sh), interpolation=cv2.INTER_AREA) + wall_small = cv2.resize(cv2.dilate(thick_mask, np.ones((3, 3), np.uint8)), + (sw, sh), interpolation=cv2.INTER_NEAREST) + else: + sam_input, wall_small = composite_rgb, thick_mask + log.append(f"SAM input {sam_input.shape[1]}x{sam_input.shape[0]} · walls at {W}x{H}") + + params = { + "seal_px": seal_px, "min_area_px": min_area_px, "max_area_frac": max_area_frac, + "enclosure_min": enclosure_min, + "n_pos": n_pos, "n_neg_wall": n_neg_wall, "n_neg_neigh": n_neg_neigh, + "neg_reach": neg_reach, "box_pad": box_pad, "use_box": bool(use_box), + "hard_clip": bool(hard_clip), "min_region_iou": min_region_iou, + "max_leak": max_leak, "leak_penalty": leak_penalty, + "fallback_region": bool(fallback_region), + "min_solidity": min_solidity, "min_extent": min_extent, + "max_vertices": max_vertices, "min_axis_frac": min_axis_frac, + "frag_max": frag_max, "merge_iou": merge_iou, "contain_max": contain_max, + } + + with sam_session() as predictor: + if predictor is None: + log.append("SAM unavailable (no torch / checkpoint). Cannot segment.") + return (cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB), walls_rgb, composite_rgb, + None, None, [], [], [], "\n".join(log), None, None, None) + rooms, prompt_viz = segment_rooms_prompted( + predictor, sam_input, wall_small, params, log) + + # Map kept rooms from the SAM small frame back to the full cropped frame. + if sf != 1.0: + for r in rooms: + seg_full = cv2.resize(r["segmentation"].astype(np.uint8), (W, H), + interpolation=cv2.INTER_NEAREST).astype(bool) + r["segmentation"] = seg_full + ys, xs = np.where(seg_full) + r["area"] = int(seg_full.sum()) + if xs.size: + r["bbox"] = (int(xs.min()), int(ys.min()), + int(xs.max() - xs.min() + 1), int(ys.max() - ys.min() + 1)) + log.append(f"[rooms] accepted after geometry + merge: {len(rooms)}") + + color = _color_overlay(cropped, rooms) + bound = _boundary_overlay(cropped, rooms) + insts = _instances(cropped, rooms) + + recs = _room_records(rooms, eff_scale, eff_dpi) + table = [[recs_r["id"], recs_r["area_px"], recs_r["area_m2"], + sum(len(p) // 2 for p in recs_r["polygons"]), + r["iou_region"], r["leak"], r["metrics"].get("solidity", ""), + int(r["metrics"].get("vertices", 0)), "region" if r["fallback"] else "sam"] + for r, recs_r in zip(rooms, recs)] + total_m2 = round(sum(x["area_m2"] for x in recs), 2) + table.append(["TOTAL", "", total_m2, "", "", "", "", "", ""]) + + h, w = cropped.shape[:2] + png, svg, js = _export_files(color, recs, (w, h)) + + log.append(f"Done · {len(rooms)} rooms · {total_m2} m² · {time.time() - t0:.1f}s") + orig = cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB) + return (orig, walls_rgb, composite_rgb, prompt_viz, color, table, [bound], insts, + "\n".join(log), png, svg, js) + + +# ───────────────────────────────────────────────────────────────────────────── +# UI +# ───────────────────────────────────────────────────────────────────────────── +with gr.Blocks(title=f"DeepPlan — Automatic Room Segmentation v{SERVICE_VERSION}") as demo: + gr.Markdown( + f"# DeepPlan — Automatic Room Segmentation (SAM) \n" + f"Device: **{_DEVICE}** ({_GPU_NAME}) · v{SERVICE_VERSION} \n" + f"Wall source: **{'wall_vectorizer_flask (vtracer, local)' if (VTRACE_WALLS and _HAS_VTRACER and _HAS_SVGPATHTOOLS) else 'UNAVAILABLE — needs vtracer + svgpathtools'}** \n" + f"Extract walls → keep only thick structural walls → find free-space regions " + f"fully enclosed by them → prompt SAM per region with positive points, " + f"negative points and a tight box → clip at the wall → validate the geometry." + ) + with gr.Row(): + with gr.Column(scale=1): + inp = gr.Image(type="pil", label="Floor plan (raster path)", height=320) + pdf_file = gr.File( + label="PDF (vector path — preferred)" if _HAS_PDF + else f"PDF vector path unavailable ({_PDF_ERR}) — pip install pymupdf", + file_types=[".pdf"], interactive=_HAS_PDF) + maxdim = gr.Slider(800, 6000, value=2000, step=100, + label="Max SAM size (px) — walls always full-res; this downscales SAM only") + with gr.Row(): + scale_denom = gr.Number( + value=100, precision=0, + label="Scale 1:N — imperial: 3/16\"=1'0\" → 64, 1/4\" → 48, 1/8\" → 96") + dpi = gr.Number(value=None, precision=0, + label="DPI (blank = 150) — room areas only; wall extraction " + "is in pixels") + btn = gr.Button("Segment rooms", variant="primary") + + with gr.Accordion("0 · PDF vector walls", open=True): + gr.Markdown( + "A vector PDF stores each wall as geometry — coordinates, stroke " + "width, fills — so there is nothing to recover by thresholding " + "pixels and **no resolution to be robust to**. Thresholds below " + "are millimetres of real building, converted through *Scale 1:N*, " + "so render DPI does not change the result. Falls back to the " + "raster path automatically if the page turns out to be a scan.") + pdf_page = gr.Number(value=1, precision=0, label="Page") + pdf_dpi = gr.Slider(72, 600, value=200, step=8, + label="Render DPI — only how finely the finished " + "geometry is drawn, not what is detected") + pdf_min_thk = gr.Slider(20, 400, value=75, step=5, + label="Min wall thickness (mm)") + pdf_max_thk = gr.Slider(100, 1200, value=500, step=10, + label="Max wall thickness (mm)") + pdf_min_len = gr.Slider(20, 3000, value=150, step=10, + label="Min wall length (mm) — low, because CAD emits " + "a wall face as many short segments which are " + "welded back together by the tracer") + pdf_strokes = gr.Checkbox(value=True, + label="A · heavy single strokes") + pdf_pairs = gr.Checkbox(value=True, + label="B · double-line walls (two parallel faces)") + pdf_fills = gr.Checkbox(value=True, + label="C · poché fills (long thin filled regions)") + + with gr.Accordion("1 · Structural walls", open=True): + ocr_erase = gr.Checkbox( + value=False, + label="OCR text-erase (needs easyocr) — also strips text-shaped wall pixels") + sam_min_stroke = gr.Slider( + 0, 40, value=10, step=1, + label="Min wall stroke (px) — the thin-line cut. Everything thinner " + "(sprinklers, electrical, plumbing, dimensions, text, symbols, " + "furniture, drafting lines) is discarded. 0 disables") + seal_px = gr.Slider(0, 30, value=6, step=1, + label="Seal gaps (px) — MORPH_CLOSE on the wall so doors " + "and wall breaks don't fuse two rooms into one region") + wall_thicken = gr.Slider(0, 12, value=SEG_WALL_THICKEN_PX, step=1, + label="Wall burn-in thickness (px) on the image SAM sees") + + with gr.Accordion("2 · Enclosure", open=True): + enclosure_min = gr.Slider( + 0.80, 1.0, value=0.98, step=0.01, + label="Min perimeter that is wall — 1.0 demands a perfectly continuous " + "boundary. Regions below this are open/incomplete and rejected") + min_area_px = gr.Slider(0, 20000, value=1500, step=100, + label="Min room area (px)") + max_area_frac = gr.Slider(0.05, 1.0, value=0.4, step=0.05, + label="Max room area (fraction of sheet)") + + with gr.Accordion("3 · Prompts", open=False): + n_pos = gr.Slider(1, 8, value=3, step=1, + label="Positive points per region (distance-transform maxima)") + n_neg_wall = gr.Slider(0, 32, value=12, step=1, + label="Negative points on the wall ring") + n_neg_neigh = gr.Slider(0, 16, value=6, step=1, + label="Negative points inside adjacent regions " + "(corridors, shafts, exterior, rooms past a door)") + neg_reach = gr.Slider(3, 200, value=60, step=1, + label="Negative reach (px) — must exceed wall thickness " + "or the adjacent-region negatives never fire") + use_box = gr.Checkbox(value=True, label="Box prompt tight around each region") + box_pad = gr.Slider(0, 20, value=2, step=1, label="Box padding (px)") + hard_clip = gr.Checkbox( + value=True, + label="Hard wall constraint — clip at the wall, then keep only the part " + "connected to the seed. Makes crossing a wall structurally impossible") + + with gr.Accordion("4 · Accept / reject", open=False): + min_region_iou = gr.Slider(0.0, 1.0, value=0.60, step=0.05, + label="Min IoU with the enclosed region") + max_leak = gr.Slider(0.0, 0.5, value=0.02, step=0.01, + label="Max boundary leakage into adjacent space") + leak_penalty = gr.Slider(0.0, 10.0, value=3.0, step=0.5, + label="Leak penalty when ranking SAM's 3 candidates " + "(objective = IoU − penalty × leak)") + fallback_region = gr.Checkbox( + value=True, + label="Keep the enclosed region when SAM's mask fails the gate — the " + "region is already proven wall-bounded") + min_solidity = gr.Slider(0.0, 1.0, value=0.55, step=0.05, + label="Min solidity (area/hull) — rect 1.0, L or T ≈0.7") + min_extent = gr.Slider(0.0, 1.0, value=0.40, step=0.05, + label="Min extent (area/bbox) — drops slivers") + max_vertices = gr.Slider(4, 64, value=16, step=1, + label="Max polygon vertices — a room is 4-12, " + "a noisy blob is dozens") + min_axis_frac = gr.Slider(0.0, 1.0, value=0.70, step=0.05, + label="Min axis-aligned outline — rooms are " + "rectilinear; pipe-following masks are not") + frag_max = gr.Slider(0.0, 0.5, value=0.05, step=0.01, + label="Max fragmentation — area outside the main part") + merge_iou = gr.Slider(0.1, 1.0, value=0.70, step=0.05, + label="Merge IoU — fuse over-segmented masks of one room") + contain_max = gr.Slider(0.5, 1.0, value=0.90, step=0.05, + label="Max containment — drop a mask this far inside " + "an already-kept one") + + with gr.Column(scale=2): + with gr.Tab("Walls"): + out_walls = gr.Image(label="Extracted walls (frontend amber overlay)", height=460) + with gr.Tab("Rooms"): + out_color = gr.Image(label="Colour-coded rooms", height=460) + out_table = gr.Dataframe( + headers=["id", "area px", "area m²", "vertices", "IoU region", + "leak", "solidity", "corners", "source"], + label="Rooms", wrap=True) + with gr.Tab("Boundaries"): + out_bound = gr.Gallery(label="Boundary overlay", height=460, columns=1) + with gr.Tab("Instances"): + out_inst = gr.Gallery(label="Individual rooms", height=460, columns=4) + with gr.Tab("Prompts"): + out_prompts = gr.Image( + label="Prompts — green = positive, red = negative, yellow = box", + height=460) + with gr.Tab("Input / composite"): + out_orig = gr.Image(label="Cropped blueprint", height=300) + out_comp = gr.Image(label="Thick-wall composite (fed to SAM)", height=300) + with gr.Row(): + dl_png = gr.File(label="PNG") + dl_svg = gr.File(label="SVG") + dl_json = gr.File(label="JSON") + out_log = gr.Textbox(label="Pipeline log", lines=10, max_lines=24) + + btn.click( + segment, + [inp, + pdf_file, pdf_page, pdf_dpi, pdf_min_thk, pdf_max_thk, pdf_min_len, + pdf_strokes, pdf_pairs, pdf_fills, + maxdim, scale_denom, dpi, + ocr_erase, sam_min_stroke, seal_px, wall_thicken, + n_pos, n_neg_wall, n_neg_neigh, neg_reach, box_pad, use_box, hard_clip, + min_area_px, max_area_frac, enclosure_min, + min_region_iou, max_leak, leak_penalty, fallback_region, + min_solidity, min_extent, max_vertices, min_axis_frac, frag_max, + merge_iou, contain_max], + [out_orig, out_walls, out_comp, out_prompts, out_color, out_table, out_bound, + out_inst, out_log, dl_png, dl_svg, dl_json], + ) + + +def _warmup() -> None: + """Load (download if missing) the SAM checkpoint. Runs in a BACKGROUND thread on + HF Spaces so the 2.4 GB download never blocks port binding (a blocked port makes + the Space look unhealthy / time out). First segment waits on the SAM lock if the + load is still running. Disable with WARMUP_SAM=0.""" + if os.environ.get("WARMUP_SAM", "1") == "0": + return + print("[startup] warming up SAM checkpoint (background)...") + with sam_session() as predictor: + print(f"[startup] SAM ready: {predictor is not None} on {_DEVICE}") + + +if __name__ == "__main__": + # Warm up in the background so demo.launch() binds the port immediately. + threading.Thread(target=_warmup, daemon=True).start() + # queue() lets long SAM jobs run without HTTP timeouts (HF proxies are strict). + demo.queue(max_size=8) + # HF Spaces / containers set GRADIO_SERVER_PORT (7860) or PORT. If neither is set, + # pass None so Gradio scans upward from 7860 for a free port. + _p = os.environ.get("GRADIO_SERVER_PORT") or os.environ.get("PORT") + demo.launch(server_name="0.0.0.0", server_port=int(_p) if _p else None)