1 Extract
rune edited this page 2026-08-23 19:30:23 +00:00
This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/usr/bin/env python3
"""
floorplan_extract.py — find the floor-plan drawing in a scanned building-case PDF
and re-draw it as normalised SVG + JSON (+ optional Graphviz DOT room graph),
so that plans from many different archives can be compared side by side.

Pipeline
--------
  1. render      page -> greyscale raster (pdftoppm)
  2. ink mask    background division + Otsu  (works for pencil, ink and blueprint)
  3. line mask   morphological opening -> only long horizontal/vertical strokes
  4. regions     cluster the line mask, score each cluster for "floor-plan-ness"
  5. vectorise   line components -> merged, snapped, orthogonal wall segments
  6. rooms       flood fill of enclosed cells -> room polygons + adjacency graph
  7. emit        <stem>.svg  <stem>.json  <stem>.dot  <stem>.preview.png

Commands
--------
  detect   file.pdf                       rank pages by floor-plan score
  extract  file.pdf --page 14 [options]   vectorise one page
  batch    indir/ -o outdir/              detect + extract for every PDF, build index

Typical use
-----------
  python3 floorplan_extract.py detect  case.pdf
  python3 floorplan_extract.py extract case.pdf --page 14 --scale 100 -o out/
  python3 floorplan_extract.py batch   ./pdfs -o ./db
  # then:  python3 floorplan_extract.py batch ... && cat db/index.jsonl

Requires: poppler-utils (pdftoppm, pdftotext), opencv-python, numpy.
"""

from __future__ import annotations

import argparse
import glob
import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
from dataclasses import dataclass, field, asdict

import cv2
import numpy as np

# --------------------------------------------------------------------------
# tunables — all lengths are fractions of the page/region size, so they are
# resolution independent
# --------------------------------------------------------------------------
CFG = dict(
    detect_dpi=110,        # cheap pass for page ranking
    extract_dpi=300,       # high quality pass for vectorising
    work_maxdim=2600,      # downscale before morphology (speed + noise)
    line_len=0.012,        # min stroke length as fraction of the sheet's long side
    frame_len=0.55,        # strokes longer than this fraction are sheet borders
    min_region=0.002,      # min cluster area as fraction of the page
    snap=0.010,            # coordinate snapping tolerance, fraction of region size
    join=0.030,            # endpoint-to-junction extension, fraction of region size
    min_wall=0.030,        # discard wall segments shorter than this
    min_room=0.004,        # min room cell area, fraction of region area
    max_room=0.400,        # max room cell area (bigger = outside/background)
)

SVG_W, SVG_H = 1000.0, 700.0    # every plan is drawn on this canvas -> comparable
SVG_PAD = 60.0

STYLE = """
  .sheet  { fill:#ffffff }
  .room   { fill:#eef2f6; stroke:none }
  .room:nth-child(even) { fill:#e7edf3 }
  .wall   { stroke:#111111; stroke-width:3.2; stroke-linecap:square }
  .thin   { stroke:#111111; stroke-width:1.4 }
  .bbox   { fill:none; stroke:#b9c4d0; stroke-width:1; stroke-dasharray:6 5 }
  .label  { font:13px 'DejaVu Sans',sans-serif; fill:#33414f; text-anchor:middle }
  .meta   { font:14px 'DejaVu Sans',sans-serif; fill:#222 }
  .metasm { font:11px 'DejaVu Sans',sans-serif; fill:#6b7783 }
  .bar    { stroke:#222; stroke-width:2 }
"""


# --------------------------------------------------------------------------
# raster helpers
# --------------------------------------------------------------------------
def page_count(pdf: str) -> int:
    out = subprocess.run(["pdfinfo", pdf], capture_output=True, text=True).stdout
    m = re.search(r"^Pages:\s+(\d+)", out, re.M)
    return int(m.group(1)) if m else 0


def render_page(pdf: str, page: int, dpi: int) -> np.ndarray:
    """Render one page to a greyscale numpy array."""
    with tempfile.TemporaryDirectory() as td:
        pre = os.path.join(td, "pg")
        subprocess.run(["pdftoppm", "-r", str(dpi), "-f", str(page), "-l", str(page),
                        "-gray", "-png", pdf, pre], check=True,
                       stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        files = sorted(glob.glob(pre + "*.png"))
        if not files:
            raise RuntimeError(f"could not render page {page} of {pdf}")
        return cv2.imread(files[0], cv2.IMREAD_GRAYSCALE)


def downscale(img: np.ndarray, maxdim: int) -> tuple[np.ndarray, float]:
    s = maxdim / max(img.shape)
    if s >= 1:
        return img, 1.0
    return cv2.resize(img, None, fx=s, fy=s, interpolation=cv2.INTER_AREA), s


def ink_mask(gray: np.ndarray) -> np.ndarray:
    """Binary ink mask, robust to paper tone, fold shadows and blueprint blue."""
    k = max(15, (min(gray.shape) // 40) | 1)
    bg = cv2.morphologyEx(gray, cv2.MORPH_CLOSE,
                          cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k)))
    norm = cv2.divide(gray, bg, scale=255)
    norm = cv2.GaussianBlur(norm, (3, 3), 0)
    _, th = cv2.threshold(norm, 0, 255,
                          cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)
    return th


def line_masks(mask: np.ndarray, L: int) -> tuple[np.ndarray, np.ndarray]:
    """Keep only strokes at least L px long, horizontally / vertically."""
    H = cv2.morphologyEx(mask, cv2.MORPH_OPEN,
                         cv2.getStructuringElement(cv2.MORPH_RECT, (L, 1)))
    V = cv2.morphologyEx(mask, cv2.MORPH_OPEN,
                         cv2.getStructuringElement(cv2.MORPH_RECT, (1, L)))
    return H, V


def strip_frame(H: np.ndarray, V: np.ndarray, shape) -> tuple[np.ndarray, np.ndarray]:
    """Drop sheet borders / fold lines / long dimension chains."""
    h, w = shape
    out = []
    for M, horiz in ((H, True), (V, False)):
        n, lab, st, _ = cv2.connectedComponentsWithStats(M, 8)
        keep = np.zeros_like(M)
        limit = CFG["frame_len"] * (w if horiz else h)
        for i in range(1, n):
            x, y, ww, hh, a = st[i]
            if (ww if horiz else hh) > limit:
                continue
            keep[lab == i] = 255
        out.append(keep)
    return out[0], out[1]


# --------------------------------------------------------------------------
# region scoring
# --------------------------------------------------------------------------
@dataclass
class Region:
    x: int
    y: int
    w: int
    h: int
    score: float = 0.0
    cells: int = 0
    fill: float = 0.0
    hv: float = 0.0

    @property
    def box(self):
        return (self.x, self.y, self.w, self.h)


def find_regions(H: np.ndarray, V: np.ndarray, page_area: int, L: int) -> list[Region]:
    lm = cv2.bitwise_or(H, V)
    d = cv2.dilate(lm, cv2.getStructuringElement(cv2.MORPH_RECT, (L, L)))
    n, lab, st, _ = cv2.connectedComponentsWithStats(d, 8)
    regs: list[Region] = []
    for i in range(1, n):
        x, y, w, h, a = st[i]
        if a < CFG["min_region"] * page_area or w < 40 or h < 40:
            continue
        sub = lm[y:y + h, x:x + w]
        kk = max(3, L // 3)
        closed = cv2.morphologyEx(sub, cv2.MORPH_CLOSE, np.ones((kk, kk), np.uint8))
        nn, _, ss, _ = cv2.connectedComponentsWithStats(255 - closed, 4)
        cells = [ss[j][4] for j in range(1, nn)
                 if CFG["min_room"] * w * h < ss[j][4] < CFG["max_room"] * w * h]
        hsum = float(H[y:y + h, x:x + w].sum()) + 1.0
        vsum = float(V[y:y + h, x:x + w].sum()) + 1.0
        hv = hsum / vsum
        fill = sum(cells) / float(w * h)
        # a floor plan = many enclosed cells, filling most of its bbox,
        # with both axes well represented (a title block is all horizontal,
        # a facade is mostly horizontal with few closed cells)
        balance = min(hv, 1 / hv)                       # 1.0 = perfectly balanced
        score = len(cells) * fill * (balance ** 0.5)
        regs.append(Region(x, y, w, h, round(score, 2), len(cells),
                           round(fill, 3), round(hv, 2)))
    regs.sort(key=lambda r: -r.score)
    return regs


def score_page(pdf: str, page: int, dpi: int) -> tuple[float, Region | None, tuple]:
    gray = render_page(pdf, page, dpi)
    small, _ = downscale(gray, CFG["work_maxdim"])
    m = ink_mask(small)
    L = max(9, int(CFG["line_len"] * max(small.shape)))
    H, V = strip_frame(*line_masks(m, L), small.shape)
    regs = find_regions(H, V, small.size, L)
    best = regs[0] if regs else None
    return (best.score if best else 0.0), best, small.shape


# --------------------------------------------------------------------------
# vectorisation
# --------------------------------------------------------------------------
@dataclass
class Seg:
    x1: float
    y1: float
    x2: float
    y2: float
    axis: str          # 'h' or 'v'
    weight: float = 1  # stroke thickness in source px (proxy for wall thickness)


def _segments_from_mask(M: np.ndarray, horiz: bool) -> list[Seg]:
    n, lab, st, _ = cv2.connectedComponentsWithStats(M, 8)
    segs = []
    for i in range(1, n):
        x, y, w, h, a = st[i]
        if horiz:
            if w < 3 or w < h:
                continue
            segs.append(Seg(x, y + h / 2.0, x + w, y + h / 2.0, "h", h))
        else:
            if h < 3 or h < w:
                continue
            segs.append(Seg(x + w / 2.0, y, x + w / 2.0, y + h, "v", w))
    return segs


def _cluster(values: list[float], tol: float) -> dict[float, float]:
    """Map each value to the mean of its tolerance cluster (1-D snapping)."""
    out = {}
    for v in sorted(set(values)):
        placed = False
        for c in list(out):
            if abs(c - v) <= tol:
                out[v] = c
                placed = True
                break
        if not placed:
            out[v] = v
    # second pass: replace representatives with cluster means
    groups: dict[float, list[float]] = {}
    for v, c in out.items():
        groups.setdefault(c, []).append(v)
    return {v: float(np.mean(groups[c])) for v, c in out.items()}


def _merge_axis(segs: list[Seg], tol: float, gap: float) -> list[Seg]:
    """Snap to shared axis coordinates and merge overlapping colinear runs."""
    if not segs:
        return []
    horiz = segs[0].axis == "h"
    key = (lambda s: s.y1) if horiz else (lambda s: s.x1)
    snap = _cluster([key(s) for s in segs], tol)
    buckets: dict[float, list[Seg]] = {}
    for s in segs:
        buckets.setdefault(snap[key(s)], []).append(s)

    merged: list[Seg] = []
    for c, group in buckets.items():
        spans = sorted(((min(s.x1, s.x2), max(s.x1, s.x2), s.weight) if horiz
                        else (min(s.y1, s.y2), max(s.y1, s.y2), s.weight)
                        for s in group))
        cur_a, cur_b, wts = spans[0][0], spans[0][1], [spans[0][2]]
        runs = []
        for a, b, wt in spans[1:]:
            if a <= cur_b + gap:
                cur_b = max(cur_b, b)
                wts.append(wt)
            else:
                runs.append((cur_a, cur_b, float(np.median(wts))))
                cur_a, cur_b, wts = a, b, [wt]
        runs.append((cur_a, cur_b, float(np.median(wts))))
        for a, b, wt in runs:
            merged.append(Seg(a, c, b, c, "h", wt) if horiz
                          else Seg(c, a, c, b, "v", wt))
    return merged


def _extend_to_junctions(segs: list[Seg], reach: float) -> list[Seg]:
    """Pull endpoints onto nearby perpendicular lines so corners actually close."""
    hs = [s for s in segs if s.axis == "h"]
    vs = [s for s in segs if s.axis == "v"]
    vx = sorted({s.x1 for s in vs})
    hy = sorted({s.y1 for s in hs})

    def nearest(val, arr):
        if not arr:
            return None
        i = int(np.argmin([abs(a - val) for a in arr]))
        return arr[i] if abs(arr[i] - val) <= reach else None

    for s in hs:
        a = nearest(s.x1, vx)
        b = nearest(s.x2, vx)
        if a is not None:
            s.x1 = a
        if b is not None:
            s.x2 = b
    for s in vs:
        a = nearest(s.y1, hy)
        b = nearest(s.y2, hy)
        if a is not None:
            s.y1 = a
        if b is not None:
            s.y2 = b
    return hs + vs


def vectorise(H: np.ndarray, V: np.ndarray, reg: Region) -> list[Seg]:
    x, y, w, h = reg.box
    subH = H[y:y + h, x:x + w]
    subV = V[y:y + h, x:x + w]
    diag = float(max(w, h))
    segs = _segments_from_mask(subH, True) + _segments_from_mask(subV, False)
    tol = CFG["snap"] * diag
    gap = CFG["join"] * diag
    segs = (_merge_axis([s for s in segs if s.axis == "h"], tol, gap)
            + _merge_axis([s for s in segs if s.axis == "v"], tol, gap))
    segs = _extend_to_junctions(segs, gap)
    minlen = CFG["min_wall"] * diag
    segs = [s for s in segs
            if max(abs(s.x2 - s.x1), abs(s.y2 - s.y1)) >= minlen]
    return segs


def rasterise_segments(segs: list[Seg], w: int, h: int, thick: int = 3) -> np.ndarray:
    canvas = np.zeros((h, w), np.uint8)
    for s in segs:
        cv2.line(canvas, (int(round(s.x1)), int(round(s.y1))),
                 (int(round(s.x2)), int(round(s.y2))), 255, thick)
    return canvas


def find_rooms(segs: list[Seg], w: int, h: int) -> list[dict]:
    """Enclosed cells of the wall graph -> room polygons."""
    walls = rasterise_segments(segs, w, h, thick=3)
    free = 255 - walls
    n, lab, st, cent = cv2.connectedComponentsWithStats(free, 4)
    rooms = []
    area_min = CFG["min_room"] * w * h
    area_max = CFG["max_room"] * w * h
    for i in range(1, n):
        x, y, ww, hh, a = st[i]
        if not (area_min < a < area_max):
            continue
        if x == 0 and y == 0 and ww == w and hh == h:
            continue                              # background
        touches_border = x <= 1 or y <= 1 or x + ww >= w - 1 or y + hh >= h - 1
        m = (lab == i).astype(np.uint8) * 255
        cnts, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        c = max(cnts, key=cv2.contourArea)
        poly = cv2.approxPolyDP(c, 0.01 * cv2.arcLength(c, True), True)
        rooms.append(dict(
            id=f"r{len(rooms) + 1}",
            label=None,                            # fill in by hand / later OCR
            area_px=int(a),
            bbox=[int(x), int(y), int(ww), int(hh)],
            centroid=[round(float(cent[i][0]), 1), round(float(cent[i][1]), 1)],
            outline=[[int(p[0][0]), int(p[0][1])] for p in poly],
            open_to_edge=bool(touches_border),
            _mask_index=i,
        ))
    return rooms, lab


def room_adjacency(rooms: list[dict], lab: np.ndarray, reach: int = 9) -> list[list[str]]:
    """Two rooms are adjacent if their cells come within `reach` px (shared wall)."""
    idx = {r["_mask_index"]: r["id"] for r in rooms}
    pairs = set()
    k = np.ones((reach, reach), np.uint8)
    for r in rooms:
        m = (lab == r["_mask_index"]).astype(np.uint8)
        grown = cv2.dilate(m, k)
        for other in np.unique(lab[grown > 0]):
            if other in idx and idx[other] != r["id"]:
                pairs.add(tuple(sorted((r["id"], idx[other]))))
    return [list(p) for p in sorted(pairs)]


# --------------------------------------------------------------------------
# metadata from the PDF text layer (best effort — these scans are OCR'd)
# --------------------------------------------------------------------------
def pdf_text(pdf: str) -> str:
    try:
        out = subprocess.run(["pdftotext", "-enc", "UTF-8", pdf, "-"],
                             capture_output=True).stdout
        return out.decode("utf-8", "replace")
    except Exception:
        return ""


def scrape_meta(text: str) -> dict:
    def first(pattern, flags=re.I):
        m = re.search(pattern, text, flags)
        return m.group(1).strip() if m else None

    return dict(
        case_no=first(r"[Bb]yggesag\s*n?r?\.?\s*[:\s]\s*([0-9]{1,4}\s*[-/]\s*[0-9]{2,4})"),
        matr_no=first(r"matr\.?\s*nr\.?\s*[:\s]\s*([0-9]{1,4}\s*[a-zæøå]{0,3})"),
        address=first(r"beliggende[:\s]+([A-ZÆØÅ][\wÆØÅæøå\. ]+\s+\d{1,3})"),
        area_m2=first(r"(\d{2,4})\s*m2?\b"),
        year=first(r"\b(19\d{2}|20\d{2})\b"),
        scale_hint=first(r"(?:MÅL|MAL|SKALA|Skala)\s*1\s*[:;]\s*(\d{2,4})"),
    )


# --------------------------------------------------------------------------
# output writers
# --------------------------------------------------------------------------
def _fit_transform(reg_w: int, reg_h: int):
    s = min((SVG_W - 2 * SVG_PAD) / reg_w, (SVG_H - 2 * SVG_PAD - 40) / reg_h)
    ox = (SVG_W - reg_w * s) / 2
    oy = (SVG_H - 40 - reg_h * s) / 2
    return s, ox, oy


def write_svg(path, segs, rooms, reg, meta, mm_per_px):
    s, ox, oy = _fit_transform(reg.w, reg.h)
    X = lambda v: round(ox + v * s, 2)
    Y = lambda v: round(oy + v * s, 2)

    parts = [
        f'<svg xmlns="http://www.w3.org/2000/svg" width="{SVG_W:.0f}" '
        f'height="{SVG_H:.0f}" viewBox="0 0 {SVG_W:.0f} {SVG_H:.0f}">',
        f"<style>{STYLE}</style>",
        "<metadata>" + json.dumps(meta, ensure_ascii=False) + "</metadata>",
        f'<rect class="sheet" x="0" y="0" width="{SVG_W:.0f}" height="{SVG_H:.0f}"/>',
        '<g id="rooms">',
    ]
    for r in rooms:
        pts = " ".join(f"{X(px)},{Y(py)}" for px, py in r["outline"])
        parts.append(f'<polygon class="room" id="{r["id"]}" points="{pts}"/>')
    parts.append("</g>")

    parts.append('<g id="walls">')
    for sg in segs:
        cls = "wall" if sg.weight >= 3 else "wall thin"
        parts.append(f'<line class="{cls}" x1="{X(sg.x1)}" y1="{Y(sg.y1)}" '
                     f'x2="{X(sg.x2)}" y2="{Y(sg.y2)}"/>')
    parts.append("</g>")

    parts.append('<g id="room-labels">')
    for r in rooms:
        cx, cy = r["centroid"]
        txt = r["label"] or r["id"]
        if mm_per_px:
            txt += f'  {r["area_px"] * (mm_per_px / 1000.0) ** 2:.1f} m²'
        parts.append(f'<text class="label" x="{X(cx)}" y="{Y(cy)}">{txt}</text>')
    parts.append("</g>")

    # footer / title block — identical on every sheet
    t = " · ".join(x for x in [
        meta.get("source"), f'p.{meta.get("page")}',
        meta.get("case_no"), meta.get("matr_no"), meta.get("address")] if x)
    parts.append(f'<text class="meta" x="{SVG_PAD}" y="{SVG_H - 34}">{t}</text>')
    sub = (f'{len(segs)} wall segments · {len(rooms)} rooms · '
           f'source bbox {reg.w}×{reg.h}px · score {reg.score}')
    if mm_per_px:
        sub += (f' · {reg.w * mm_per_px / 1000:.1f}×{reg.h * mm_per_px / 1000:.1f} m'
                f' @1:{meta.get("scale")}')
    parts.append(f'<text class="metasm" x="{SVG_PAD}" y="{SVG_H - 16}">{sub}</text>')

    if mm_per_px:                                   # 1 m scale bar
        L = 1000.0 / mm_per_px * s
        x0, y0 = SVG_W - SVG_PAD - L, SVG_H - 26
        parts.append(f'<line class="bar" x1="{x0:.1f}" y1="{y0}" '
                     f'x2="{x0 + L:.1f}" y2="{y0}"/>')
        parts.append(f'<text class="metasm" x="{x0 + L / 2:.1f}" y="{y0 - 6}" '
                     f'text-anchor="middle">1 m</text>')

    parts.append("</svg>")
    with open(path, "w", encoding="utf-8") as f:
        f.write("\n".join(parts))


def write_dot(path, rooms, edges, meta):
    lines = ['graph plan {', '  graph [layout=neato, overlap=false, splines=true];',
             '  node  [shape=box, style="rounded,filled", fillcolor="#eef2f6", '
             'fontname="Helvetica", fontsize=10];',
             f'  label="{meta.get("source")} p.{meta.get("page")}";']
    for r in rooms:
        cx, cy = r["centroid"]
        lab = r["label"] or r["id"]
        if r.get("area_m2"):
            lab += f'\\n{r["area_m2"]} m²'
        lines.append(f'  {r["id"]} [label="{lab}", pos="{cx / 40:.2f},{-cy / 40:.2f}!"];')
    for a, b in edges:
        lines.append(f"  {a} -- {b};")
    lines.append("}")
    with open(path, "w", encoding="utf-8") as f:
        f.write("\n".join(lines))


def write_preview(path, gray, reg, segs):
    vis = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
    x, y,.