| |
| """Read the circuits off an electrical floor plan with onnxruntime alone (no torch). |
| |
| pip install onnxruntime numpy pillow |
| python predict.py plan.png [--symbol-px 16] [--out circuits.png] [--json circuits.json] |
| |
| The model was trained on sheets where a duplex receptacle is ~12-24 px across. --symbol-px is the size of a |
| receptacle circle in YOUR image; the sheet is rescaled so symbols land in that range before inference. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import pathlib |
|
|
| import numpy as np |
| import onnxruntime as ort |
| from PIL import Image, ImageDraw |
|
|
| import decode as D |
|
|
| HERE = pathlib.Path(__file__).resolve().parent |
| TILE, OVERLAP = 512, 64 |
| TRAIN_SYMBOL_PX = 16.0 |
|
|
|
|
| def run_sheet(sess, gray: np.ndarray): |
| """Tile a grayscale sheet (0..255, HxW) through the model and stitch the stride-2 outputs by centre crop.""" |
| H, W = gray.shape |
| gh, gw = (H + 1) // 2, (W + 1) // 2 |
| peaks = np.zeros((len(D.CLASSES), gh, gw), np.float32) |
| size = np.zeros((2, gh, gw), np.float32) |
| wire = np.zeros((gh, gw), np.float32) |
| step = TILE - 2 * OVERLAP |
| for y0 in range(-OVERLAP, max(1, H - OVERLAP), step): |
| for x0 in range(-OVERLAP, max(1, W - OVERLAP), step): |
| tile = np.full((TILE, TILE), 255, np.float32) |
| ys, xs = max(0, y0), max(0, x0) |
| ye, xe = min(H, y0 + TILE), min(W, x0 + TILE) |
| tile[ys - y0:ye - y0, xs - x0:xe - x0] = gray[ys:ye, xs:xe] |
| p, s, w = sess.run(None, {"image": tile[None, None]}) |
| |
| cy0, cx0 = (y0 + OVERLAP) // 2, (x0 + OVERLAP) // 2 |
| n = step // 2 |
| o = OVERLAP // 2 |
| ya, xa = max(0, cy0), max(0, cx0) |
| yb, xb = min(gh, cy0 + n), min(gw, cx0 + n) |
| if yb <= ya or xb <= xa: |
| continue |
| sl = (slice(o + ya - cy0, o + yb - cy0), slice(o + xa - cx0, o + xb - cx0)) |
| peaks[:, ya:yb, xa:xb] = p[0][:, sl[0], sl[1]] |
| size[:, ya:yb, xa:xb] = s[0][:, sl[0], sl[1]] |
| wire[ya:yb, xa:xb] = w[0, 0][sl] |
| return peaks, size, wire |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("image") |
| ap.add_argument("--model", default=str(HERE / "circuits.onnx")) |
| 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") |
| a = ap.parse_args() |
| img = Image.open(a.image).convert("L") |
| 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) |
| inv = 1 / scale |
| for d in res["devices"]: |
| d["box"] = [round(v * inv, 1) for v in d["box"]] |
| d["class"] = D.cname(d["cls"]) |
| for c in res["circuits"]: |
| c["length_px"] = round(c["length_px"] * inv, 1) |
| for r in res["runs"]: |
| r.pop("points", None) |
| for k, c in enumerate(res["circuits"]): |
| kinds = [res["devices"][i]["class"] for i in c["devices"]] |
| print(f"circuit {k + 1}: {len(c['devices'])} devices ({', '.join(sorted(set(kinds)))}), " |
| f"home run: {'yes' if c['homeruns'] else 'no'}, ~{c['length_px']:.0f} px of wiring") |
| if res["unwired"]: |
| print(f"{len(res['unwired'])} powered symbols with no wiring found") |
| if a.json: |
| 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) |
|
|
|
|
| PALETTE = [(31, 119, 180), (214, 39, 40), (44, 160, 44), (148, 103, 189), (255, 127, 14), (23, 190, 207), |
| (227, 119, 194), (140, 86, 75)] |
|
|
|
|
| def draw_result(img, res): |
| d = ImageDraw.Draw(img) |
| cid = {i: k for k, c in enumerate(res["circuits"]) for i in c["devices"]} |
| for dv in res["devices"]: |
| col = PALETTE[cid[dv["id"]] % len(PALETTE)] if dv["id"] in cid else (130, 130, 130) |
| d.rectangle(dv["box"], outline=col, width=3) |
| if dv["id"] in cid: |
| d.text((dv["box"][0], dv["box"][1] - 11), str(cid[dv["id"]] + 1), fill=col) |
| return img |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
