Spaces:
Paused
Paused
| """ | |
| 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. | |
| # βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ | |
| # PDF input. A PDF is RENDERED TO AN IMAGE and then goes through exactly the same | |
| # wall extraction as an uploaded image β one code path, one set of thresholds, and | |
| # no dependence on a sibling module, which is what makes it work on a Space where | |
| # only this file is deployed. Needs pymupdf (`pip install pymupdf`); nothing else. | |
| try: | |
| import fitz # type: ignore # pymupdf | |
| _HAS_FITZ = True | |
| _FITZ_ERR = "" | |
| except Exception as _pexc: | |
| fitz = None # type: ignore | |
| _HAS_FITZ = False | |
| _FITZ_ERR = f"{type(_pexc).__name__}: {_pexc}" | |
| print(f"[pdf] render-to-image: {'available' if _HAS_FITZ else 'OFF (' + _FITZ_ERR + ')'}") | |
| PDF_MAX_DIM = int(os.environ.get("PDF_MAX_DIM", "12000")) # cap the rendered page's long edge | |
| _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 <path>, 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" <path {_attrs_str(by_el[k][0], ' '.join(by_el[k][1]))}/>" for k in order] | |
| hdr = [f'<svg xmlns="{_SVG_NS}"'] | |
| if view_box: | |
| hdr.append(f' viewBox="{view_box}"') | |
| if width: | |
| hdr.append(f' width="{width}"') | |
| if height: | |
| hdr.append(f' height="{height}"') | |
| hdr.append(">") | |
| svg = "".join(hdr) + "\n" + "\n".join(kept) + "\n</svg>" | |
| 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 | |
| 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'<svg xmlns="http://www.w3.org/2000/svg" width="{w}" height="{h}" ' | |
| f'viewBox="0 0 {w} {h}">'] | |
| 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'<polygon points="{pts}" fill="{fill}" fill-opacity="0.45" ' | |
| f'stroke="{fill}" stroke-width="2"/>') | |
| parts.append(f'<text x="{r["centroid"][0]}" y="{r["centroid"][1]}" ' | |
| f'font-size="14" fill="#000">{r["id"]}</text>') | |
| parts.append("</svg>") | |
| 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 _pdf_to_image(pdf_file: Any, page_no: int, render_dpi: int, | |
| log: List[str]) -> Optional[np.ndarray]: | |
| """Render one PDF page to a BGR image. | |
| The page becomes an ordinary raster and then goes through the identical wall | |
| extraction an uploaded image would β the PDF is an input format here, not a | |
| separate pipeline. Returns None (having logged why) if it cannot render. | |
| """ | |
| if pdf_file is None: | |
| return None | |
| if not _HAS_FITZ: | |
| log.append(f"[pdf] cannot render β pymupdf missing ({_FITZ_ERR}); " | |
| f"pip install pymupdf, or upload the sheet as an image") | |
| return None | |
| path = getattr(pdf_file, "name", None) or str(pdf_file) | |
| try: | |
| doc = fitz.open(path) | |
| except Exception as exc: | |
| log.append(f"[pdf] could not open ({type(exc).__name__}: {exc})") | |
| return None | |
| try: | |
| if doc.page_count < 1: | |
| log.append("[pdf] document has no pages") | |
| 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} " | |
| f"of {doc.page_count}") | |
| page = doc[idx] | |
| # Cap the render so a big sheet at a high DPI cannot exhaust memory before | |
| # anything is extracted. The cap also honours EXTRACT_MAX_DIM: extraction | |
| # downscales to it regardless, so rendering larger only costs memory β | |
| # an A1 at 600 dpi is a 1.2 GB array that is then thrown away. Raise | |
| # EXTRACT_MAX_DIM if you want the finer render to actually be used. | |
| cap = min(PDF_MAX_DIM, EXTRACT_MAX_DIM) if EXTRACT_MAX_DIM else PDF_MAX_DIM | |
| dpi = max(36, int(render_dpi)) | |
| rect = page.rect | |
| long_pt = max(rect.width, rect.height) | |
| if long_pt > 0 and long_pt * dpi / 72.0 > cap: | |
| dpi = max(36, int(cap * 72.0 / long_pt)) | |
| log.append(f"[pdf] {int(render_dpi)} dpi would exceed the {cap}px working " | |
| f"size β rendering at {dpi} dpi instead") | |
| pix = page.get_pixmap(matrix=fitz.Matrix(dpi / 72.0, dpi / 72.0), alpha=False) | |
| img = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, pix.n) | |
| if pix.n == 4: | |
| bgr = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR) | |
| elif pix.n == 3: | |
| bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) | |
| else: | |
| bgr = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) | |
| log.append(f"[pdf] page {idx + 1}/{doc.page_count} rendered at {dpi} dpi β " | |
| f"{bgr.shape[1]}x{bgr.shape[0]} px; extracting walls from the image") | |
| return bgr | |
| except Exception as exc: | |
| log.append(f"[pdf] render failed ({type(exc).__name__}: {exc})") | |
| return None | |
| finally: | |
| try: | |
| doc.close() | |
| except Exception: | |
| pass | |
| def segment(image: Optional[Image.Image], | |
| pdf_file: Any, pdf_page: int, pdf_dpi: int, | |
| 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.", | |
| None, None, None) | |
| t0 = time.time() | |
| log: List[str] = [] | |
| # A PDF is rendered to an image and then treated exactly like an uploaded | |
| # image β same crop, same wall extraction, same thresholds. The PDF is an | |
| # input format, not a second pipeline. | |
| bgr: Optional[np.ndarray] = None | |
| if pdf_file is not None: | |
| bgr = _pdf_to_image(pdf_file, pdf_page, pdf_dpi, log) | |
| if bgr is None and image is None: | |
| return (None, None, None, None, None, [], [], [], | |
| "\n".join(log + ["Could not render the PDF, and no image was uploaded."]), | |
| None, None, None) | |
| if bgr is not None and image is not None: | |
| log.append("[pdf] both a PDF and an image were supplied β using the PDF") | |
| if bgr is None: | |
| 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) | |
| elif EXTRACT_MAX_DIM and max(bgr.shape[:2]) > EXTRACT_MAX_DIM: | |
| s = EXTRACT_MAX_DIM / float(max(bgr.shape[:2])) | |
| bgr = cv2.resize(bgr, (max(1, int(bgr.shape[1] * s)), max(1, int(bgr.shape[0] * s))), | |
| interpolation=cv2.INTER_AREA) | |
| log.append(f"[walls] rendered page > {EXTRACT_MAX_DIM}px β downscaled for memory") | |
| 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 (image)", height=320) | |
| pdf_file = gr.File( | |
| label="PDF β rendered to an image, then walls extracted from it" | |
| if _HAS_FITZ | |
| else f"PDF unavailable ({_FITZ_ERR}) β pip install pymupdf", | |
| file_types=[".pdf"], interactive=_HAS_FITZ) | |
| 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 input", open=True): | |
| gr.Markdown( | |
| "A PDF page is **rendered to an image**, and the walls are then " | |
| "extracted from that image by exactly the same method as a " | |
| "direct image upload. Render DPI is the one setting that " | |
| "matters: it decides how many pixels the sheet gets, and a wall " | |
| "drawn as two thin lines needs enough of them to stay two lines. " | |
| "300 dpi on an A1 sheet is a good default; below ~200 the faces " | |
| "blur together.") | |
| pdf_page = gr.Number(value=1, precision=0, label="Page") | |
| pdf_dpi = gr.Slider(72, 600, value=300, step=8, | |
| label="Render DPI β how many pixels the page becomes") | |
| 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, | |
| 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) | |