← back to Wallco Ai

scripts/clean-edge-upscale.py

262 lines

#!/usr/bin/env python3
r"""
clean-edge-upscale.py — clean-edge / 150-DPI upscaler for flat-color wallco designs.

WHY: ComfyUI renders at 1024px. A 36" tile at 1024px is only ~28 DPI, so every
flat-color silhouette edge stair-steps / aliases at real print size (Steve flagged
#53677 2026-05-27: "edges must be clean and not jagged. increase dpi if needed").
Print standard = 150 DPI -> a 36" tile is 5400x5400px.

WHAT: vector-traces the design with potrace (resolution-independent razor edges),
snaps the result back to its solid inks, and rasterizes at 5400px. If the trace is
degenerate (lost a real ink, or the design isn't flat-color), it falls back to a
plain LANCZOS upscale so the design is still produced — never made worse.

This is the ORCHESTRATOR. The heavy lifting (dominant-ink quantize, potrace per
layer, ink-snap) lives in scripts/vectorize/vectorize.py — this wraps it with:
  - design-id -> source-PNG resolution (via data/designs.json filename + IMG_DIR)
  - auto ink-count detection (don't invent ghost inks on a clean 2-tone)
  - degenerate-trace detection -> LANCZOS fallback
  - clean_<id>_<ts>.png output naming into data/generated/
  - Round-1 SACRED: never overwrites the source PNG

NOT wired into the publish flow or the catalog — build/validate tool only.

============================ USAGE ============================
  # by design id (resolves data/designs.json -> data/generated/<filename>)
  python3 scripts/clean-edge-upscale.py 53677

  # by explicit path
  python3 scripts/clean-edge-upscale.py data/generated/foo.png --id foo

  # tuning / options
  python3 scripts/clean-edge-upscale.py 53677 \
      --target 5400 \           # output px (5400 = 150 DPI @ 36")
      --max-inks 6 \            # cap on auto-detected solid inks
      --min-ink-share 0.8 \     # % population an ink needs to count as real
      --turdsize 2 \            # potrace: drop traced specks smaller than this
      --alphamax 1.0 \          # potrace: corner smoothing (0=polygon,1.33=round)
      --out /tmp/clean_53677.png \
      --force-lanczos \         # skip the trace, LANCZOS only (for comparison)
      --snap-fallback \         # snap the LANCZOS fallback back to inks too
      --keep-aa \               # keep cairo's 1px edge AA (skip the ink-snap)
      --json                    # machine-readable summary on stdout

  # runs under any python3 — auto re-execs under scripts/vectorize/.venv if the
  # vector deps (cairosvg/numpy/PIL) aren't importable in the current interpreter.
===============================================================
"""
import argparse
import json
import os
import sys
import time
from pathlib import Path

# --- venv bootstrap: the vector engine needs cairosvg/numpy/PIL from the venv ---
_HERE = Path(__file__).resolve().parent
_VECDIR = _HERE / "vectorize"
_VENV_PY = _VECDIR / ".venv" / "bin" / "python"
try:
    import numpy as np  # noqa: F401
    import cairosvg  # noqa: F401
    from PIL import Image  # noqa: F401
except Exception:
    if _VENV_PY.exists() and os.environ.get("_CEU_REEXEC") != "1":
        os.environ["_CEU_REEXEC"] = "1"
        os.execv(str(_VENV_PY), [str(_VENV_PY), str(Path(__file__).resolve()), *sys.argv[1:]])
    raise

import numpy as np
from PIL import Image

# the trace engine lives next door; import its functions directly
sys.path.insert(0, str(_VECDIR))
from vectorize import quantize_dominant, build_svg, snap_to_inks  # noqa: E402

PROJ = _HERE.parent
IMG_DIR = PROJ / "data" / "generated"
DESIGNS_JSON = PROJ / "data" / "designs.json"


# ------------------------------- helpers ----------------------------------- #
def resolve_source(arg: str):
    """Return (src_path, design_id). arg may be a file path or a design id."""
    p = Path(arg)
    if p.exists() and p.is_file():
        return p, p.stem
    # treat as design id -> look up filename in data/designs.json
    if not DESIGNS_JSON.exists():
        sys.exit(f"[clean-edge] '{arg}' is not a file and {DESIGNS_JSON} is missing")
    data = json.loads(DESIGNS_JSON.read_text())
    rows = data if isinstance(data, list) else (data.get("designs") or data.get("rows") or [])
    want = str(arg)
    row = next((r for r in rows if str(r.get("id")) == want or str(r.get("dig_number")) == want), None)
    if not row:
        sys.exit(f"[clean-edge] design id {arg} not found in {DESIGNS_JSON.name}")
    fn = row.get("filename")
    if not fn:
        sys.exit(f"[clean-edge] design {arg} has no filename column")
    src = IMG_DIR / Path(fn).name  # basename only (PII-free, matches localPathForDesign)
    if not src.exists():
        sys.exit(f"[clean-edge] source PNG not found: {src}")
    return src, want


def auto_ink_count(img: Image.Image, max_inks: int, min_share: float, merge_dist: float = 24.0):
    """Count real solid inks: dominant colors >merge_dist apart that each hold
    >= min_share % of the pixels. Avoids inventing ghost inks from AA fringe on a
    clean 2-tone, and avoids forcing extra inks on a 5-ink design."""
    arr = np.array(img.convert("RGB")).reshape(-1, 3)
    uniq, counts = np.unique(arr, axis=0, return_counts=True)
    order = np.argsort(-counts)
    total = arr.shape[0]
    inks, shares = [], []
    for i in order:
        c = uniq[i].astype(np.int32)
        if all(np.sqrt(((c - np.array(k, np.int32)) ** 2).sum()) > merge_dist for k in inks):
            inks.append(tuple(int(x) for x in c))
            shares.append(counts[i] / total)
        if len(inks) >= max_inks:
            break
    n = sum(1 for s in shares if s * 100.0 >= min_share)
    return max(2, n)


def ink_shares(labels, n):
    tot = labels.size
    return [float((labels == i).sum()) / tot for i in range(n)]


def render_shares(arr, colors):
    """Population share of each ink in a snapped raster (exact color match)."""
    pal = np.array(colors, np.int32)
    flat = arr.reshape(-1, 3).astype(np.int32)
    tot = flat.shape[0]
    out = []
    for c in pal:
        out.append(float((flat == c).all(1).sum()) / tot)
    return out


def aa_ramp_fraction(arr, colors, tol=20):
    """Fraction of pixels NOT within tol of any ink = gradient/AA-blur ramp.
    Lower = crisper, more solid screen-print. (Same metric as vectorize/compare.py.)"""
    pal = np.array(colors, np.int32)
    flat = arr.reshape(-1, 3).astype(np.int32)
    out = np.empty(flat.shape[0])
    step = 2_000_000
    for s in range(0, flat.shape[0], step):
        ch = flat[s:s + step]
        d = np.sqrt(((ch[:, None, :] - pal[None, :, :]) ** 2).sum(2)).min(1)
        out[s:s + step] = d
    return float((out > tol).mean())


# -------------------------------- main ------------------------------------- #
def main():
    ap = argparse.ArgumentParser(description="clean-edge / 150-DPI upscaler (potrace, LANCZOS fallback)")
    ap.add_argument("input", help="design id (looked up in designs.json) OR a path to a PNG")
    ap.add_argument("--id", default=None, help="override the id used in the output filename")
    ap.add_argument("--target", type=int, default=5400, help="output px square (5400 = 150 DPI @ 36in)")
    ap.add_argument("--max-inks", type=int, default=6)
    ap.add_argument("--min-ink-share", type=float, default=0.8, help="%% population for an ink to count")
    ap.add_argument("--turdsize", type=int, default=2)
    ap.add_argument("--alphamax", type=float, default=1.0)
    ap.add_argument("--out", default=None, help="explicit output path (default data/generated/clean_<id>_<ts>.png)")
    ap.add_argument("--force-lanczos", action="store_true", help="skip the trace, LANCZOS only")
    ap.add_argument("--snap-fallback", action="store_true", help="snap the LANCZOS fallback back to the inks")
    ap.add_argument("--keep-aa", action="store_true", help="skip the nearest-ink snap on the vector raster")
    ap.add_argument("--json", action="store_true", help="emit a machine-readable summary on stdout")
    args = ap.parse_args()

    src, did = resolve_source(args.input)
    did = args.id or did
    ts = int(time.time() * 1000)
    out = Path(args.out) if args.out else (IMG_DIR / f"clean_{did}_{ts}.png")
    if out.resolve() == src.resolve():
        sys.exit("[clean-edge] refusing to overwrite the source PNG (Round-1 sacred)")
    out.parent.mkdir(parents=True, exist_ok=True)

    img = Image.open(src).convert("RGB")
    n = auto_ink_count(img, args.max_inks, args.min_ink_share)
    summary = {"src": str(src), "id": did, "out": str(out), "target": args.target,
               "src_size": list(img.size), "auto_inks": n}

    method = "potrace"
    reason = ""
    if args.force_lanczos:
        method, reason = "lanczos", "forced"

    if method == "potrace":
        try:
            labels, colors = quantize_dominant(img, n)
            used = sorted({int(x) for x in np.unique(labels)})
            remap = {old: new for new, old in enumerate(used)}
            labels = np.vectorize(remap.get)(labels).astype(np.int32)
            colors = [colors[o] for o in used]
            src_shares = ink_shares(labels, len(colors))

            import tempfile
            with tempfile.TemporaryDirectory() as td:
                svg, colors = build_svg(labels, colors, args.turdsize, args.alphamax, td)
            n_paths = svg.count("<path")
            png_bytes = cairosvg.svg2png(bytestring=svg.encode(),
                                         output_width=args.target, output_height=args.target)
            tmp = str(out) + ".aa.tmp.png"
            Path(tmp).write_bytes(png_bytes)
            arr = np.array(Image.open(tmp).convert("RGB"))
            os.remove(tmp)
            if not args.keep_aa:
                arr = snap_to_inks(arr, colors)
            out_shares = render_shares(arr, colors)

            # --- degeneracy detection ---
            # (1) no traced geometry at all -> just the base rect = degenerate
            if n_paths == 0:
                raise RuntimeError("trace produced no paths (solid fill only)")
            # (2) a real ink (>=2% of source) collapsed in the output (<25% of its
            #     source share AND <0.5% absolute) -> the subject/ink was lost.
            #     NOTE: thin-line FRAGMENTATION is acceptable (ink area is preserved,
            #     just broken into blobs) — that does NOT trip this. Only true LOSS does.
            for i, (ss, os_) in enumerate(zip(src_shares, out_shares)):
                if ss >= 0.02 and os_ < ss * 0.25 and os_ < 0.005:
                    raise RuntimeError(
                        f"ink #{'%02x%02x%02x' % colors[i]} collapsed "
                        f"{ss*100:.1f}%% -> {os_*100:.2f}%% (lost)")

            colors_out = len(np.unique(arr.reshape(-1, 3), axis=0))
            Image.fromarray(arr).save(out)
            summary.update({"method": "potrace", "inks": ["#%02x%02x%02x" % c for c in colors],
                            "n_paths": n_paths, "out_colors": int(colors_out),
                            "aa_ramp_pct": round(aa_ramp_fraction(arr, colors) * 100, 4)})
        except Exception as e:
            method, reason = "lanczos", f"degenerate trace: {e}"

    if method == "lanczos":
        arr = np.array(img.resize((args.target, args.target), Image.LANCZOS).convert("RGB"))
        if args.snap_fallback:
            labels, colors = quantize_dominant(img, n)
            arr = snap_to_inks(arr, colors)
            summary["aa_ramp_pct"] = round(aa_ramp_fraction(arr, colors) * 100, 4)
        colors_out = len(np.unique(arr.reshape(-1, 3), axis=0))
        Image.fromarray(arr).save(out)
        summary.update({"method": "lanczos", "fallback_reason": reason, "out_colors": int(colors_out)})

    summary["method"] = method
    if args.json:
        print(json.dumps(summary))
    else:
        print(f"[clean-edge] {src.name} ({img.size[0]}px) -> {method} -> {out.name} "
              f"@ {args.target}px")
        if method == "potrace":
            print(f"[clean-edge]   inks={summary['inks']} paths={summary['n_paths']} "
                  f"out_colors={summary['out_colors']} aa_ramp={summary['aa_ramp_pct']}%")
        else:
            print(f"[clean-edge]   LANCZOS fallback ({reason}); out_colors={summary['out_colors']}")
    return summary


if __name__ == "__main__":
    main()