gendesign/backend/scripts/spike_plan_vectorize.py

"""Spike — contour-vectorize floor-plan rasters (PNG/JPG → SVG) via Potrace.

Forgejo issue #299 (Phase 1 exploration SPIKE). **This is a throwaway probe, NOT
production wiring**: it answers one question — is Potrace contour tracing good
enough to turn floor-plan rasters into compact, readable SVG for the planning UI
catalog and PDF reports? No ML, no room semantics, no DOM.RF coupling.

WHAT IT MEASURES
----------------
For each input raster the pipeline runs:

  1. load (Pillow) → grayscale (``L``)
  2. binarise: pixels below ``--threshold`` (0-255) become ink (black), the rest
     become paper (white). Optional ``--invert`` for light-on-dark plans.
  3. write a 1-bit PBM bitmap (Potrace's native input) to the output dir
  4. ``potrace -s`` → SVG (centerline-free, filled-contour tracing)
  5. record size metrics: raster bytes, svg bytes, ratio (raster / svg)

It also (optionally, ``--render-back``) rasterises the produced SVG back to PNG
via ``rsvg-convert`` so a human can eyeball SVG-vs-original side by side and
judge whether wall/room contours survived the round trip.

WHY POTRACE IS SOURCE-AGNOSTIC
------------------------------
Potrace traces the boundary between ink and paper regions of a bitmap. It does
not care where the raster came from — a Wikimedia line-drawing plan, a scanned
blueprint, or a DOM.RF UI tile all reduce to "dark strokes on a light field"
after thresholding. So the Phase 1 question ("are the contours traceable and how
much do they compress?") is validly answered on any representative floor-plan
rasters; see the spike doc for the explicit DOM.RF caveat.

USAGE
-----
    # vectorise every raster in a folder, collect a metrics table
    uv run python backend/scripts/spike_plan_vectorize.py \
        --in-dir /tmp/plan-spike/in --out-dir /tmp/plan-spike/out

    # also render the SVGs back to PNG for visual QA
    uv run python backend/scripts/spike_plan_vectorize.py \
        --in-dir /tmp/plan-spike/in --out-dir /tmp/plan-spike/out --render-back

Requires the ``potrace`` binary on PATH (``brew install potrace``) and, for
``--render-back``, ``rsvg-convert`` (``brew install librsvg``). Pillow ships with
the backend deps.
"""

from __future__ import annotations

import argparse
import logging
import shutil
import statistics
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path

from PIL import Image

logger = logging.getLogger("spike_plan_vectorize")

_RASTER_SUFFIXES = {".png", ".jpg", ".jpeg", ".bmp", ".gif", ".tif", ".tiff"}
_POTRACE_TIMEOUT_S = 120
_RSVG_TIMEOUT_S = 120


@dataclass(frozen=True)
class VectorizeResult:
    """One raster's trip through the pipeline."""

    name: str
    raster_bytes: int
    svg_bytes: int
    svg_path: Path
    rendered_path: Path | None

    @property
    def ratio(self) -> float:
        """Compression ratio raster / svg (>1 means SVG is smaller)."""
        return self.raster_bytes / self.svg_bytes if self.svg_bytes else 0.0

    def metrics_line(self) -> str:
        return (
            f"{self.name:<44} "
            f"raster={self.raster_bytes:>9}B  "
            f"svg={self.svg_bytes:>8}B  "
            f"ratio={self.ratio:>6.2f}x"
        )


def _require_binary(name: str) -> str:
    """Resolve an external binary on PATH or exit with a clear message."""
    path = shutil.which(name)
    if path is None:
        logger.error("required binary %r not found on PATH", name)
        raise SystemExit(
            f"{name!r} not found — install it (e.g. `brew install {name}`) and retry"
        )
    return path


def raster_to_pbm(src: Path, pbm_path: Path, *, threshold: int, invert: bool) -> None:
    """Load a raster, grayscale + threshold it, and write a 1-bit PBM bitmap.

    Potrace consumes 1-bit bitmaps (PBM/PGM/BMP). We binarise with a fixed
    threshold so the spike's behaviour is deterministic and explainable — black
    (ink) is what Potrace traces, white is background. ``invert`` flips the test
    for light-stroke-on-dark plans.
    """
    with Image.open(src) as im:
        gray = im.convert("L")
    # point(): pixel < threshold → ink (0), else paper (255). Pillow's "1" mode
    # then packs to a true 1-bit bitmap that Potrace reads natively.
    cutoff = threshold
    if invert:
        bitmap = gray.point(lambda px: 0 if px >= cutoff else 255).convert("1")
    else:
        bitmap = gray.point(lambda px: 0 if px < cutoff else 255).convert("1")
    bitmap.save(pbm_path)


def pbm_to_svg(potrace_bin: str, pbm_path: Path, svg_path: Path) -> None:
    """Trace a PBM bitmap to SVG with ``potrace -s`` (SVG backend)."""
    try:
        subprocess.run(
            [potrace_bin, "-s", "-o", str(svg_path), str(pbm_path)],
            check=True,
            capture_output=True,
            timeout=_POTRACE_TIMEOUT_S,
        )
    except subprocess.CalledProcessError as exc:
        stderr = exc.stderr.decode("utf-8", "replace").strip()
        logger.error("potrace failed for %s: %s", pbm_path.name, stderr)
        raise


def svg_to_png(rsvg_bin: str, svg_path: Path, png_path: Path, *, width: int) -> None:
    """Render an SVG back to PNG via ``rsvg-convert`` for visual QA.

    We composite onto an explicit white background. Potrace's SVG is black
    filled paths on a *transparent* canvas; without ``--background-color=white``
    rsvg renders the black fill onto transparency and a flattening viewer shows
    a solid black tile. White matches the real catalog/PDF use case anyway.
    """
    try:
        subprocess.run(
            [
                rsvg_bin, "-w", str(width),
                "--background-color=white",
                "-o", str(png_path), str(svg_path),
            ],
            check=True,
            capture_output=True,
            timeout=_RSVG_TIMEOUT_S,
        )
    except subprocess.CalledProcessError as exc:
        stderr = exc.stderr.decode("utf-8", "replace").strip()
        logger.error("rsvg-convert failed for %s: %s", svg_path.name, stderr)
        raise


def vectorize_one(
    src: Path,
    out_dir: Path,
    *,
    potrace_bin: str,
    rsvg_bin: str | None,
    threshold: int,
    invert: bool,
    render_width: int,
) -> VectorizeResult:
    """Run the full PNG/JPG → PBM → SVG (→ PNG) pipeline for a single raster."""
    stem = src.stem
    pbm_path = out_dir / f"{stem}.pbm"
    svg_path = out_dir / f"{stem}.svg"

    raster_to_pbm(src, pbm_path, threshold=threshold, invert=invert)
    pbm_to_svg(potrace_bin, pbm_path, svg_path)

    rendered_path: Path | None = None
    if rsvg_bin is not None:
        rendered_path = out_dir / f"{stem}.rendered.png"
        svg_to_png(rsvg_bin, svg_path, rendered_path, width=render_width)

    return VectorizeResult(
        name=src.name,
        raster_bytes=src.stat().st_size,
        svg_bytes=svg_path.stat().st_size,
        svg_path=svg_path,
        rendered_path=rendered_path,
    )


def discover_rasters(in_dir: Path) -> list[Path]:
    """Return sorted raster files in ``in_dir`` (non-recursive)."""
    return sorted(
        p for p in in_dir.iterdir() if p.is_file() and p.suffix.lower() in _RASTER_SUFFIXES
    )


def summarise(results: list[VectorizeResult]) -> None:
    """Print a per-file metrics table plus min/median/max compression."""
    if not results:
        logger.warning("no results to summarise")
        return

    print("\n=== per-raster metrics ===")
    for r in results:
        print(r.metrics_line())

    ratios = sorted(r.ratio for r in results)
    print("\n=== compression summary ===")
    print(f"samples       : {len(ratios)}")
    print(f"min ratio     : {min(ratios):.2f}x")
    print(f"median ratio  : {statistics.median(ratios):.2f}x")
    print(f"max ratio     : {max(ratios):.2f}x")
    total_raster = sum(r.raster_bytes for r in results)
    total_svg = sum(r.svg_bytes for r in results)
    print(f"aggregate     : {total_raster}B raster -> {total_svg}B svg "
          f"({total_raster / total_svg:.2f}x overall)")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Spike: contour-vectorize floor-plan rasters (PNG/JPG → SVG) via Potrace.",
    )
    parser.add_argument(
        "--in-dir", type=Path, required=True,
        help="folder containing input rasters (PNG/JPG/...)",
    )
    parser.add_argument(
        "--out-dir", type=Path, required=True,
        help="folder for output .pbm/.svg (and .rendered.png with --render-back)",
    )
    parser.add_argument(
        "--threshold", type=int, default=128,
        help="grayscale ink/paper cutoff 0-255 (default 128)",
    )
    parser.add_argument(
        "--invert", action="store_true",
        help="treat light strokes on a dark field as ink",
    )
    parser.add_argument(
        "--render-back", action="store_true",
        help="rasterise each SVG back to PNG via rsvg-convert for visual QA",
    )
    parser.add_argument(
        "--render-width", type=int, default=900,
        help="width in px for --render-back output (default 900)",
    )
    return parser


def main(argv: list[str] | None = None) -> int:
    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s")
    args = build_parser().parse_args(argv)

    in_dir: Path = args.in_dir
    out_dir: Path = args.out_dir
    if not in_dir.is_dir():
        logger.error("input dir does not exist: %s", in_dir)
        return 2
    if not 0 <= args.threshold <= 255:
        logger.error("--threshold must be 0-255, got %d", args.threshold)
        return 2
    out_dir.mkdir(parents=True, exist_ok=True)

    potrace_bin = _require_binary("potrace")
    rsvg_bin = _require_binary("rsvg-convert") if args.render_back else None

    rasters = discover_rasters(in_dir)
    if not rasters:
        logger.error("no raster files found in %s", in_dir)
        return 1
    logger.info("found %d raster(s) in %s", len(rasters), in_dir)

    results: list[VectorizeResult] = []
    failed = 0
    for src in rasters:
        try:
            results.append(
                vectorize_one(
                    src,
                    out_dir,
                    potrace_bin=potrace_bin,
                    rsvg_bin=rsvg_bin,
                    threshold=args.threshold,
                    invert=args.invert,
                    render_width=args.render_width,
                )
            )
            logger.info("vectorized %s", src.name)
        except Exception:
            # Log + continue: one bad raster must not abort the batch, but we
            # never swallow silently — the traceback is recorded and the file
            # is counted as a failure in the final tally.
            failed += 1
            logger.exception("failed to vectorize %s", src.name)

    summarise(results)
    if failed:
        logger.warning("%d/%d raster(s) failed", failed, len(rasters))
    return 0 if results else 1


if __name__ == "__main__":
    sys.exit(main())