Decoder update: predict.py

predict.py

@@ -11,6 +11,7 @@ from __future__ import annotations
 
 import argparse
 import json
+import math
 import pathlib
 
 import numpy as np
@@ -54,6 +55,232 @@ def run_sheet(sess, gray: np.ndarray):
     return peaks, size, wire
 
 
+def decode_two_scale(peaks, size, wire, peaks2, size2, k, cfg=None):
+    """decode() plus home runs whose arrowheads only fire on a coarser pass of the same sheet (arrows drawn large
+    next to small fixtures, e.g. Windsor E103). peaks2/size2: the model's outputs on the sheet shrunk by 1/k (so
+    k = this pass's px per coarse-pass px, i.e. cfg["arrow_scale"]). Wiring is the normal pass's, untouched: the
+    coarse arrows only add home runs, each landing on a device that pass already found."""
+    cfg = D.CFG if cfg is None else cfg
+    base = D.decode(peaks, size, wire, cfg)
+    dets = D.detect(peaks, size, cfg)
+    fine = [d for d in dets if d["cls"] == D.ARROW]
+    extra = []
+    for a in D.detect(peaks2, size2, cfg):
+        if a["cls"] != D.ARROW:
+            continue
+        cx, cy = a["cx"] * k, a["cy"] * k
+        if any(math.hypot(cx - f["cx"], cy - f["cy"]) < cfg.get("arrow_merge", 8) for f in fine + extra):
+            continue
+        extra.append({**a, "cx": cx, "cy": cy, "box": [v * k for v in a["box"]]})
+    if not extra:
+        return base
+    merged = [dict(d) for d in dets + extra]
+    for i, d in enumerate(merged):
+        d["id"] = i
+    res = D.decode(peaks, size, wire, cfg, dets=merged)
+
+    cen = lambda b: ((b[0] + b[2]) / 2, (b[1] + b[3]) / 2)
+    bdev = [d for d in base["devices"] if d["cls"] != D.ARROW]
+    rbox = {d["id"]: d["box"] for d in res["devices"]}
+    has_hr = {h["device"] for h in base["homeruns"]}
+    circ_of = {i: c for c in base["circuits"] for i in c["devices"]}
+    for h in res["homeruns"]:
+        if h["arrow"] < len(dets):
+            continue  # a normal-pass arrow: base already decided that home run
+        p = cen(rbox[h["device"]])
+        m = min(bdev, key=lambda d: (math.dist(cen(d["box"]), p), d["id"]), default=None)
+        if m is None or math.dist(cen(m["box"]), p) > 4 or m["id"] in has_hr:
+            continue
+        a = res["devices"][h["arrow"]]
+        aid = len(base["devices"])
+        base["devices"].append({"id": aid, "cls": D.ARROW, "score": a["score"], "box": a["box"]})
+        base["homeruns"].append({"device": m["id"], "arrow": aid})
+        base["runs"].append({"devices": [m["id"]], "homerun": True, "arrow": aid, "length_px": 0.0, "points": []})
+        has_hr.add(m["id"])
+        c = circ_of.get(m["id"])
+        if c is None:
+            c = {"devices": [m["id"]], "homeruns": [], "length_px": 0.0}
+            base["circuits"].append(c)
+            circ_of[m["id"]] = c
+        c["homeruns"] = sorted(c["homeruns"] + [aid])
+        if m["id"] in base["unwired"]:
+            base["unwired"].remove(m["id"])
+    base["homeruns"].sort(key=lambda h: h["device"])
+    return base
+
+
+def bridge_dashes(peaks, size, wire, cfg=None):
+    """Close the gaps of dashed runs (UK switching links) in the wire map before decode(). Without this every dash
+    is its own stroke: its two free ends become unrecognised devices and a dashed run decodes as a string of
+    dash-length fragments. Two free skeleton ends (not on a detected symbol) are joined when they are at most
+    cfg["dash_gap"] grid px apart and each one's outward direction, measured cfg["dash_look"] px back along its
+    stroke, points at the other within cos >= cfg["dash_cos"]; closest pairs first, each end used once.
+    dash_gap 0 (the default) returns the map unchanged."""
+    cfg = D.CFG if cfg is None else cfg
+    gap = cfg.get("dash_gap", 0)
+    if gap <= 0:
+        return wire
+    cos_min, look = cfg.get("dash_cos", 0.75), int(cfg.get("dash_look", 4))
+    dets = [d for d in D.detect(peaks, size, cfg) if d["cls"] != D.ARROW]
+    sk = D.zhang_suen((wire > cfg["wire_thr"]).astype(np.uint8))
+    H, W = sk.shape
+    pts = set(zip(*(a.tolist() for a in np.nonzero(sk))))
+    ends = []
+    for y, x in sorted(pts):
+        t, n = D.crossing(sk, y, x, H, W)
+        if not (t == 1 and n <= 2):
+            continue
+        if any(D.box_dist(y, x, d["box"]) <= cfg["attach"] for d in dets):
+            continue  # lands on a symbol: decode() attaches it
+        path, cur = [(y, x)], (y, x)
+        for _ in range(look):
+            nxt = [(cur[0] + dy, cur[1] + dx) for dy, dx in D.NB if (cur[0] + dy, cur[1] + dx) in pts
+                   and (cur[0] + dy, cur[1] + dx) not in path]
+            if not nxt:
+                break
+            cur = nxt[0]
+            path.append(cur)
+        vy, vx = y - cur[0], x - cur[1]
+        nrm = math.sqrt(vy * vy + vx * vx)  # sqrt, not hypot: bit-identical to circuit-model.js
+        if nrm < 1.5:
+            continue  # a speck: no direction to continue along
+        ends.append((y, x, vy / nrm, vx / nrm))
+    cell = {}
+    for i, (y, x, _, _) in enumerate(ends):
+        cell.setdefault((int(y // gap), int(x // gap)), []).append(i)
+    cand = []
+    for i, (y, x, uy, ux) in enumerate(ends):
+        cy, cx = int(y // gap), int(x // gap)
+        for oy in (-1, 0, 1):
+            for ox in (-1, 0, 1):
+                for j in cell.get((cy + oy, cx + ox), []):
+                    if j <= i:
+                        continue
+                    y2, x2, uy2, ux2 = ends[j]
+                    d = math.sqrt((y2 - y) * (y2 - y) + (x2 - x) * (x2 - x))
+                    if d == 0 or d > gap:
+                        continue
+                    if (uy * (y2 - y) + ux * (x2 - x)) / d >= cos_min and (uy2 * (y - y2) + ux2 * (x - x2)) / d >= cos_min:
+                        cand.append((d, i, j))
+    out, used = wire.copy(), set()
+    for d, i, j in sorted(cand):
+        if i in used or j in used:
+            continue
+        used.update((i, j))
+        (y1, x1), (y2, x2) = ends[i][:2], ends[j][:2]
+        n = math.ceil(d) * 2
+        for k in range(n + 1):
+            t = k / n
+            # floor(v + 0.5), not round(): Python rounds halves to even, JS does not
+            out[math.floor(y1 + (y2 - y1) * t + 0.5), math.floor(x1 + (x2 - x1) * t + 0.5)] = np.float32(1.0)
+    return out
+
+
+def estimate_skew(img: Image.Image, max_deg: float = 4.0, work_px: int = 900) -> float:
+    """Degrees to rotate the sheet so its drawn lines sit square. Plan sheets are dominated by horizontal and vertical
+    strokes, so the right angle is the one where ink piles into the fewest rows (largest row-sum variance).
+    Coarse 0.5 degree sweep, then 0.1; numpy and PIL only."""
+    g = img.convert("L")
+    k = work_px / max(g.size)
+    if k < 1:
+        g = g.resize((max(1, round(g.width * k)), max(1, round(g.height * k))), Image.BILINEAR)
+    ink = Image.fromarray(((np.asarray(g, np.float32) < 140) * 255).astype(np.uint8))
+
+    def score(deg):
+        rows = (np.asarray(ink.rotate(deg, resample=Image.BILINEAR, fillcolor=0), np.float32) / 255).sum(axis=1)
+        return float(np.var(rows))
+    best = max(np.arange(-max_deg, max_deg + 1e-9, 0.5), key=score)
+    best = max(np.arange(best - 0.5, best + 0.5 + 1e-9, 0.1), key=score)
+    # only trust a tilt that clearly sharpens the row profile over the unrotated sheet
+    return float(best) if score(best) > 1.05 * score(0.0) else 0.0
+
+
+def deskew(img: Image.Image, min_deg: float = 0.25):
+    """(square image, angle applied). Tilts under min_deg are left alone: rotating resamples every stroke."""
+    deg = estimate_skew(img)
+    if abs(deg) < min_deg:
+        return img, 0.0
+    return img.rotate(deg, resample=Image.BICUBIC, fillcolor=255), deg
+
+
+SCALE_CANDIDATES = (10, 12, 14, 16, 18, 20, 23, 26, 30, 35)       # receptacle diameter in page px
+
+
+def _ink_windows(gray: np.ndarray, win: int, n: int = 3):
+    """Top-left corners of the ``n`` most ink-dense, non-overlapping ``win`` x ``win`` windows of a sheet."""
+    H, W = gray.shape
+    if H <= win or W <= win:
+        return [(0, 0)]
+    b = max(8, win // 4)
+    ink = (gray < 128).astype(np.float32)
+    hh, ww = H // b, W // b
+    cells = ink[: hh * b, : ww * b].reshape(hh, b, ww, b).mean(axis=(1, 3))
+    k = max(1, win // b)
+    from scipy.ndimage import uniform_filter
+    dens = uniform_filter(cells, size=k, mode="constant")
+    out = []
+    for _ in range(n):
+        i = int(np.argmax(dens))
+        cy, cx = divmod(i, ww)
+        if dens[cy, cx] <= 0:
+            break
+        y0 = int(np.clip(cy * b - win // 2, 0, H - win))
+        x0 = int(np.clip(cx * b - win // 2, 0, W - win))
+        out.append((y0, x0))
+        dens[max(0, cy - k):cy + k + 1, max(0, cx - k):cx + k + 1] = -1
+    return out or [(0, 0)]
+
+
+def scale_evidence(res):
+    """Self-consistency of one decode, needing no labels: at the right symbol scale the model finds confident symbols
+    AND the wiring it traces ends on them. Returns (score, parts)."""
+    arrow = D.CLASSES.index("homerun_arrow")
+    devs = [d for d in res["devices"] if d["cls"] not in (arrow, D.VIRTUAL)]       # real symbols the model found
+    scores = sorted((d.get("score", 0.0) for d in devs), reverse=True)[:25]
+    conf = float(np.mean(scores)) if scores else 0.0
+    runs = res["runs"]
+    real = set(d["id"] for d in devs)
+    ends = sum(len(r["devices"]) for r in runs)
+    landed = sum(1 for r in runs for i in r["devices"] if i in real)
+    snapped = landed / ends if ends else 0.0
+    return conf * (0.25 + 0.75 * snapped) * min(1.0, len(devs) / 6.0), {"conf": round(conf, 3), "snapped": round(snapped, 3), "devices": len(devs), "runs": len(runs)}
+
+
+def estimate_symbol_px(sess, img: Image.Image, candidates=SCALE_CANDIDATES, window_model_px=512, windows=3, cfg=None):
+    """Pick the page's symbol size (receptacle diameter in page px) without asking: decode a few ink-dense crops at every
+    candidate scale and keep the scale whose decode is most self-consistent (``scale_evidence``). Returns
+    (best_symbol_px, {candidate: score})."""
+    gray = np.asarray(img.convert("L"), np.float32)
+    cfg = {**(D.CFG if cfg is None else cfg), "arrow_scale": 1.0}
+    table = {}
+    for c in candidates:
+        k = TRAIN_SYMBOL_PX / c                                # page px -> model px
+        win = int(round(window_model_px / k))
+        vals = []
+        for y0, x0 in _ink_windows(gray, win, windows):
+            crop = Image.fromarray(gray[y0:y0 + win, x0:x0 + win].astype(np.uint8))
+            crop = crop.resize((max(32, round(crop.width * k)), max(32, round(crop.height * k))), Image.LANCZOS)
+            vals.append(scale_evidence(decode_sheet(sess, crop, cfg))[0])
+        table[c] = round(float(np.mean(vals)), 4) if vals else 0.0
+    best = max(table, key=table.get)
+    return best, table
+
+
+def decode_sheet(sess, img: Image.Image, cfg=None):
+    """Model + decoder on a sheet already at the model's symbol size. With cfg["arrow_scale"] > 1, a second pass on
+    the sheet shrunk by that factor adds home runs whose arrowheads only fire there (decode_two_scale). With
+    cfg["dash_gap"] > 0, dashed runs are closed first (bridge_dashes)."""
+    cfg = D.CFG if cfg is None else cfg
+    peaks, size, wire = run_sheet(sess, np.asarray(img, np.float32))
+    wire = bridge_dashes(peaks, size, wire, cfg)
+    f = cfg.get("arrow_scale", 1.0)
+    if f <= 1.0:
+        return D.decode(peaks, size, wire, cfg)
+    small = img.resize((max(1, round(img.width / f)), max(1, round(img.height / f))), Image.LANCZOS)
+    peaks2, size2, _ = run_sheet(sess, np.asarray(small, np.float32))
+    return decode_two_scale(peaks, size, wire, peaks2, size2, img.width / small.width, cfg)
+
+
 def main():
     ap = argparse.ArgumentParser()
     ap.add_argument("image")
@@ -61,14 +288,20 @@ def main():
     ap.add_argument("--symbol-px", type=float, default=20, help="receptacle circle diameter in the input (20 suits a 150-dpi PDF render)")
     ap.add_argument("--out")
     ap.add_argument("--json")
+    ap.add_argument("--no-deskew", action="store_true", help="skip straightening a tilted scan or photo (on by default; tilts under 0.25 degrees are left alone)")
+    ap.add_argument("--arrow-scale", type=float, default=None,
+                    help="also read home-run arrows from the sheet shrunk by this factor (default: decoder_config.json arrow_scale; 1 = off)")
     a = ap.parse_args()
-    img = Image.open(a.image).convert("L")
+    shown = Image.open(a.image).convert("L")
+    img, tilt = (shown, 0.0) if a.no_deskew else deskew(shown)
+    if tilt:
+        print(f"straightened the sheet by {tilt:+.1f} degrees")
+        shown = img
     scale = TRAIN_SYMBOL_PX / a.symbol_px
     if abs(scale - 1) > 0.05:
         img = img.resize((round(img.width * scale), round(img.height * scale)), Image.LANCZOS)
     sess = ort.InferenceSession(a.model, providers=["CPUExecutionProvider"])
-    peaks, size, wire = run_sheet(sess, np.asarray(img, np.float32))
-    res = D.decode(peaks, size, wire)
+    res = decode_sheet(sess, img, D.CFG if a.arrow_scale is None else {**D.CFG, "arrow_scale": a.arrow_scale})
     inv = 1 / scale
     for d in res["devices"]:
         d["box"] = [round(v * inv, 1) for v in d["box"]]
@@ -84,9 +317,10 @@ def main():
     if res["unwired"]:
         print(f"{len(res['unwired'])} powered symbols with no wiring found")
     if a.json:
+        res["deskew_degrees"] = round(tilt, 2)
         pathlib.Path(a.json).write_text(json.dumps(res, indent=1))
     if a.out:
-        draw_result(Image.open(a.image).convert("RGB"), res).save(a.out)
+        draw_result(shown.convert("RGB"), res).save(a.out)
 
 
 PALETTE = [(31, 119, 180), (214, 39, 40), (44, 160, 44), (148, 103, 189), (255, 127, 14), (23, 190, 207),