Update app.py

app.py

@@ -66,17 +66,21 @@ import gradio as gr
 # room_validation, mask_to_polygon. Falls back to the inline port only if import
 # fails. Force the port with USE_DEEPPLAN_WALLS=0.
 # ─────────────────────────────────────────────────────────────────────────────
-# Vector wall extraction from PDF. Optional (needs pymupdf); absent → the PDF tab
-# reports why and the raster path is used.
+# 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 pdf_walls  # type: ignore
-    _HAS_PDF = pdf_walls.HAS_FITZ
-    _PDF_ERR = "" if _HAS_PDF else pdf_walls.FITZ_ERROR
+    import fitz  # type: ignore  # pymupdf
+    _HAS_FITZ = True
+    _FITZ_ERR = ""
 except Exception as _pexc:
-    pdf_walls = None  # type: ignore
-    _HAS_PDF = False
-    _PDF_ERR = f"{type(_pexc).__name__}: {_pexc}"
-print(f"[pdf] vector wall extraction: {'available' if _HAS_PDF else 'OFF (' + _PDF_ERR + ')'}")
+    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
@@ -2102,67 +2106,73 @@ def _export_files(overlay_rgb: np.ndarray, recs: List[Dict[str, Any]],
 # ─────────────────────────────────────────────────────────────────────────────
 # Gradio callback
 # ─────────────────────────────────────────────────────────────────────────────
-def _walls_from_pdf(pdf_file: Any, page_no: int, render_dpi: int, scale_denom: int,
-                    min_thk_mm: float, max_thk_mm: float, min_wall_len_mm: float,
-                    use_strokes: bool, use_pairs: bool, use_fills: bool,
-                    log: List[str]) -> Optional[Tuple[np.ndarray, np.ndarray, int]]:
-    """Vector path: PDF page → (cropped BGR, wall mask, effective dpi).
-
-    Returns None when it cannot apply, having said why in the log, so the caller
-    falls back to the raster pipeline instead of failing.
+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_PDF:
-        log.append(f"[pdf] vector extraction unavailable ({_PDF_ERR}) — using raster path")
+    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:
-        with open(path, "rb") as fh:
-            data = fh.read()
-        doc = pdf_walls.open_document(data)
+        doc = fitz.open(path)
     except Exception as exc:
-        log.append(f"[pdf] {exc} — using raster path")
+        log.append(f"[pdf] could not open ({type(exc).__name__}: {exc})")
         return None
-
-    idx = max(0, min(int(page_no) - 1, doc.page_count - 1))
-    if idx != int(page_no) - 1:
-        log.append(f"[pdf] page {page_no} out of range; using page {idx + 1} of {doc.page_count}")
-    page = doc[idx]
-
     try:
-        mask, info, plog = pdf_walls.extract_wall_mask(
-            page, dpi=int(render_dpi), scale_denominator=int(scale_denom),
-            min_thickness_mm=float(min_thk_mm), max_thickness_mm=float(max_thk_mm),
-            min_wall_length_mm=float(min_wall_len_mm),
-            use_strokes=bool(use_strokes), use_pairs=bool(use_pairs),
-            use_fills=bool(use_fills))
-        log += plog
-        if not info.get("vector") or info.get("wall_px", 0) == 0:
-            log.append("[pdf] falling back to the raster pipeline")
+        if doc.page_count < 1:
+            log.append("[pdf] document has no pages")
             return None
-        page_bgr = pdf_walls.render_page(page, int(render_dpi))
+        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] vector extraction failed ({type(exc).__name__}: {exc}) — using raster path")
+        log.append(f"[pdf] render failed ({type(exc).__name__}: {exc})")
         return None
-
-    # Same title-block crop the raster path applies, derived from the rendered
-    # page so the mask and the image stay in one coordinate space.
-    cropped, m2 = stage2_crop_drawing(page_bgr)
-    log.append(m2)
-    mask = mask[:cropped.shape[0], :cropped.shape[1]]
-    if mask.shape[:2] != cropped.shape[:2]:
-        mask = cv2.resize(mask, (cropped.shape[1], cropped.shape[0]),
-                          interpolation=cv2.INTER_NEAREST)
-    log.append("[pdf] source: VECTOR geometry — thresholds are physical (mm), "
-               "not pixels, so render DPI does not affect the result")
-    return cropped, mask, int(render_dpi)
+    finally:
+        try:
+            doc.close()
+        except Exception:
+            pass
 
 
 def segment(image: Optional[Image.Image],
             pdf_file: Any, pdf_page: int, pdf_dpi: int,
-            pdf_min_thk: float, pdf_max_thk: float, pdf_min_len: float,
-            pdf_strokes: bool, pdf_pairs: bool, pdf_fills: bool,
             maxdim: int, scale_denom: int, dpi: int,
             ocr_erase: bool, sam_min_stroke: int, seal_px: int, wall_thicken: int,
             n_pos: int, n_neg_wall: int, n_neg_neigh: int, neg_reach: int,
@@ -2175,43 +2185,35 @@ def segment(image: Optional[Image.Image],
             merge_iou: float, contain_max: float):
     if image is None and pdf_file is None:
         return (None, None, None, None, None, [], [], [],
-                "Upload a floor plan image, or a PDF for the vector path.",
+                "Upload a floor plan image or a PDF.",
                 None, None, None)
     t0 = time.time()
     log: List[str] = []
 
-    # Vector first when a PDF is supplied. A vector PDF carries the wall geometry
-    # exactly — coordinates, stroke widths, fills — so there is nothing to recover
-    # by thresholding pixels, and no render DPI for the thresholds to depend on.
-    vec = _walls_from_pdf(pdf_file, pdf_page, pdf_dpi, scale_denom,
-                          pdf_min_thk, pdf_max_thk, pdf_min_len,
-                          pdf_strokes, pdf_pairs, pdf_fills, log)
-    if vec is not None:
-        cropped, wall_mask, eff_dpi = vec
-        eff_scale = int(scale_denom) if scale_denom else 100
-        return _segment_from_walls(cropped, wall_mask, eff_dpi, eff_scale, t0, log,
-                                   maxdim, sam_min_stroke, seal_px, wall_thicken,
-                                   n_pos, n_neg_wall, n_neg_neigh, neg_reach,
-                                   box_pad, use_box, hard_clip,
-                                   min_area_px, max_area_frac, enclosure_min,
-                                   min_region_iou, max_leak, leak_penalty,
-                                   fallback_region, min_solidity, min_extent,
-                                   max_vertices, min_axis_frac, frag_max,
-                                   merge_iou, contain_max)
-    if image is None:
-        return (None, None, None, None, None, [], [], [],
-                "\n".join(log + ["No usable vector geometry and no image uploaded."]),
-                None, None, None)
+    # 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")
 
-    # Wall extraction in this app is VECTOR — NOT the heuristic Stages 2-4. The
-    # blueprint is cropped to drop the title block, then traced to vector walls
-    # locally by vtracer; the downstream thick-stroke filter keeps structural walls
-    # only. Downscale only as a memory guard for oversized sheets.
-    src_pil = image
-    if EXTRACT_MAX_DIM and max(image.size) > EXTRACT_MAX_DIM:
-        src_pil = _downscale_pil(image.convert("RGB"), EXTRACT_MAX_DIM)
-        log.append(f"[walls] image > {EXTRACT_MAX_DIM}px — downscaled for memory")
-    bgr = cv2.cvtColor(np.array(src_pil.convert("RGB")), cv2.COLOR_RGB2BGR)
+    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².
@@ -2352,11 +2354,12 @@ with gr.Blocks(title=f"DeepPlan — Automatic Room Segmentation v{SERVICE_VERSIO
     )
     with gr.Row():
         with gr.Column(scale=1):
-            inp = gr.Image(type="pil", label="Floor plan (raster path)", height=320)
+            inp = gr.Image(type="pil", label="Floor plan (image)", height=320)
             pdf_file = gr.File(
-                label="PDF (vector path — preferred)" if _HAS_PDF
-                      else f"PDF vector path unavailable ({_PDF_ERR}) — pip install pymupdf",
-                file_types=[".pdf"], interactive=_HAS_PDF)
+                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():
@@ -2368,32 +2371,18 @@ with gr.Blocks(title=f"DeepPlan — Automatic Room Segmentation v{SERVICE_VERSIO
                                       "is in pixels")
             btn = gr.Button("Segment rooms", variant="primary")
 
-            with gr.Accordion("0 · PDF vector walls", open=True):
+            with gr.Accordion("0 · PDF input", open=True):
                 gr.Markdown(
-                    "A vector PDF stores each wall as geometry — coordinates, stroke "
-                    "width, fills — so there is nothing to recover by thresholding "
-                    "pixels and **no resolution to be robust to**. Thresholds below "
-                    "are millimetres of real building, converted through *Scale 1:N*, "
-                    "so render DPI does not change the result. Falls back to the "
-                    "raster path automatically if the page turns out to be a scan.")
+                    "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=200, step=8,
-                                    label="Render DPI — only how finely the finished "
-                                          "geometry is drawn, not what is detected")
-                pdf_min_thk = gr.Slider(20, 400, value=75, step=5,
-                                        label="Min wall thickness (mm)")
-                pdf_max_thk = gr.Slider(100, 1200, value=500, step=10,
-                                        label="Max wall thickness (mm)")
-                pdf_min_len = gr.Slider(20, 3000, value=150, step=10,
-                                        label="Min wall length (mm) — low, because CAD emits "
-                                              "a wall face as many short segments which are "
-                                              "welded back together by the tracer")
-                pdf_strokes = gr.Checkbox(value=True,
-                                          label="A · heavy single strokes")
-                pdf_pairs = gr.Checkbox(value=True,
-                                        label="B · double-line walls (two parallel faces)")
-                pdf_fills = gr.Checkbox(value=True,
-                                        label="C · poché fills (long thin filled regions)")
+                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(
@@ -2497,8 +2486,7 @@ with gr.Blocks(title=f"DeepPlan — Automatic Room Segmentation v{SERVICE_VERSIO
     btn.click(
         segment,
         [inp,
-         pdf_file, pdf_page, pdf_dpi, pdf_min_thk, pdf_max_thk, pdf_min_len,
-         pdf_strokes, pdf_pairs, pdf_fills,
+         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,