AI Flow: Reverse Engineering

#!/usr/bin/env python3
"""Draw a sample of pad-to-pad connections and render them for checking by eye.

    connect_sample.py <reverse-flow folder> <out folder> <name> --per 75 [--exclude other.json ...]

A connection is an edge of the shortest spanning tree over the pins of a group in detect/connect-graph.json: the pad-to-pad
joins the copper is claimed to make, nearest first, so paths stay short enough to look along. Strata: big groups (at least
BIG pins) and small ones. The order comes from sha256(name + hash of the graph), so nobody chooses it, and a sample drawn
from the same file is the same sample. Writes sample.json and one image per connection: the photo left, the claimed copper
tinted right. Top row: both faces of the whole path. Bottom row: each pad's edge up close. Cards carry no net names, so a checker is not led by them.
"""
import sys, json, math, hashlib, random, struct, os, argparse
import numpy as np, cv2
BIG = 8
ap = argparse.ArgumentParser(); ap.add_argument("flow"); ap.add_argument("out"); ap.add_argument("name"); ap.add_argument("--per", type=int, default=75); ap.add_argument("--exclude", nargs="*", default=[]); ap.add_argument("--controls", type=int, default=0)
a = ap.parse_args(); R = a.flow.rstrip("/") + "/"; os.makedirs(a.out, exist_ok=True)
raw = open(R + "detect/connect-graph.json", "rb").read(); g = json.loads(raw)
ms = json.load(open(R + "detect/attrib/measure.json")); at = json.load(open(R + "detect/attribution.json")); cop = json.load(open(R + "copper.json"))
pads = {}
for p in ms["pads"]: pads.setdefault(f"{p['part']}.{p['num']}", p)
via_pos = {v["id"]: v["pos_mm"] for v in cop["vias"]}
def read_rle(p):
    b = open(p, "rb").read(); w, h = struct.unpack("<II", b[:8]); r = np.frombuffer(b[8:], dtype="<u4").reshape(-1, 2)
    return np.repeat(r[:, 0], r[:, 1]).reshape(h, w)
maps = {s: read_rle(R + f"detect/attrib/pieces-{s}.rle") for s in ("top", "bottom")}
imgs = {s: cv2.imread(R + f"board/{s}.jpg") for s in ("top", "bottom")}
piece_rows = at["pieces"]
contacts, solid = {}, {}
for c in g["contacts"]:
    contacts.setdefault(c["pin"], []).append(c)
    if c["tier"] == 1: solid.setdefault(c["pin"], []).append(c)
links = {}
for l in g["links"]:
    if l["tier"] <= 3:
        x, y = l["pieces"]; links.setdefault(x, []).append((y, l)); links.setdefault(y, []).append((x, l))
conns = []
for gr in g["groups"]:
    pins = [p for p in gr["pins"] if p in pads and not p.startswith("PHOTO1.") and p in solid]
    if len(pins) < 2: continue
    # Prim over Euclidean distance
    done = [pins[0]]; rest = set(pins[1:])
    while rest:
        best = min(((math.dist((pads[u]["x"], pads[u]["y"]), (pads[v]["x"], pads[v]["y"])), u, v) for u in done for v in sorted(rest)))
        conns.append({"a": best[1], "b": best[2], "group": gr["group"], "net": gr["net"], "groupPins": len(gr["pins"]), "lengthMm": round(best[0], 2)}); done.append(best[2]); rest.discard(best[2])
# When the spanning edges are too few for the sample, the other pairs in the same groups are claims too: a group says all
# its pins are joined. Those are added, nearest first in each group, and drawn in the same order.
tree = {f"{min(c['a'], c['b'])}~{max(c['a'], c['b'])}" for c in conns}
for gr in g["groups"]:
    pins = [p for p in gr["pins"] if p in pads and not p.startswith("PHOTO1.") and p in solid]
    for i, u in enumerate(pins):
        for v in pins[i + 1:]:
            if u.split(".")[0] == v.split(".")[0]: continue
            if f"{min(u, v)}~{max(u, v)}" in tree: continue
            conns.append({"a": u, "b": v, "group": gr["group"], "net": gr["net"], "groupPins": len(gr["pins"]), "lengthMm": round(math.dist((pads[u]["x"], pads[u]["y"]), (pads[v]["x"], pads[v]["y"])), 2), "pair": True})
# Pads whose edge did not settle it still lie on copper: those pairs are claims of a lower grade, reported as their own grade.
seen_pairs = {f"{min(c['a'], c['b'])}~{max(c['a'], c['b'])}" for c in conns}
for c in conns: c["grade"] = "solid"
node_pins = {}
for pin, cs in contacts.items():
    if pin.startswith("PHOTO1.") or pin not in pads: continue
    for c in cs: node_pins.setdefault(c["node"], set()).add(pin)
gname = {n: gr for gr in g["groups"] for n in gr["nodes"]}
for n, pins in node_pins.items():
    pins = sorted(pins)
    for i, u in enumerate(pins):
        for v in pins[i + 1:]:
            k = f"{min(u, v)}~{max(u, v)}"
            if k in seen_pairs or u.split(".")[0] == v.split(".")[0]: continue
            seen_pairs.add(k); gr = gname[n]
            conns.append({"a": u, "b": v, "group": gr["group"], "net": gr["net"], "groupPins": len(pins), "lengthMm": round(math.dist((pads[u]["x"], pads[u]["y"]), (pads[v]["x"], pads[v]["y"])), 2), "pair": True, "grade": "unsettled"})
excl = set()
for f in a.exclude: excl |= {c["id"] for c in json.load(open(f))["connections"]}
for c in conns: c["id"] = f"{min(c['a'], c['b'])}~{max(c['a'], c['b'])}"
pop = [c for c in conns if c["id"] not in excl]
seed = hashlib.sha256((a.name + hashlib.sha256(raw).hexdigest()).encode()).hexdigest()
rng = random.Random(seed); pop.sort(key=lambda c: c["id"]); rng.shuffle(pop)
pop.sort(key=lambda c: (0 if c["grade"] == "solid" else 1))  # every solid claim is taken before any unsettled one
big = [c for c in pop if c["groupPins"] >= BIG][:a.per]; small = [c for c in pop if c["groupPins"] < BIG][:a.per]
sample = big + small
# Controls for the checker: the two ends of a resistor or capacitor, each on its own copper, offered as one connection. Correct answer: not joined.
ctrl = []
for part in sorted({p.split(".")[0] for p in contacts if p[0] in "RC" and p.split(".")[0][1:].isdigit()}):
    u, v = f"{part}.1", f"{part}.2"
    if u in contacts and v in contacts and {x["node"] for x in contacts[u]}.isdisjoint({x["node"] for x in contacts[v]}):
        ctrl.append({"a": u, "b": v, "group": -1, "net": "", "groupPins": 0, "lengthMm": 0, "grade": "control", "control": True, "nodes": [contacts[u][0]["node"], contacts[v][0]["node"]], "id": f"control:{part}"})
rng.shuffle(ctrl); sample = sample + ctrl[:a.controls]; rng.shuffle(sample)
def path(c):
    """Nodes and vias from pad a to pad b, fewest hops."""
    start = {x["node"] for x in contacts[c["a"]]}; goal = {x["node"] for x in contacts[c["b"]]}
    prev = {n: None for n in start}; q = list(start); hit = None
    while q:
        u = q.pop(0)
        if u in goal: hit = u; break
        for v, l in links.get(u, []):
            if v not in prev: prev[v] = (u, l); q.append(v)
    nodes, vias = [], []
    while hit is not None:
        nodes.append(hit); p = prev[hit]
        if p: vias.append(p[1]["ids"][0]); hit = p[0]
        else: hit = None
    return nodes, vias
T = 250
def tile(s, x0, y0, x1, y1, ids, marks, label):
    """Photo and the same photo with the claimed copper tinted, side by side, T px each."""
    px = lambda v: int(round((v - 100) * 40))
    H, W = maps[s].shape; b = (max(px(x0), 0), max(px(y0), 0), min(px(x1), W), min(px(y1), H))
    im = imgs[s][b[1]:b[3], b[0]:b[2]]; lab = maps[s][b[1]:b[3], b[0]:b[2]]
    if im.size == 0: return np.zeros((T, 2 * T, 3), np.uint8)
    mk = np.isin(lab, list(ids)) if ids else np.zeros(lab.shape, bool)
    f = T / max(im.shape[:2]); L = cv2.resize(im, None, fx=f, fy=f, interpolation=cv2.INTER_CUBIC)
    over = im.copy(); over[mk] = (0.55 * over[mk] + 0.45 * np.array([0, 200, 255])).astype(np.uint8)
    O = cv2.resize(over, None, fx=f, fy=f, interpolation=cv2.INTER_NEAREST)
    for name, q, hw, hh in marks:
        u, v = ((q[0] - 100) * 40 - b[0]) * f, ((q[1] - 100) * 40 - b[1]) * f
        if name == "v": cv2.circle(O, (int(u), int(v)), 8, (255, 255, 0), 1); continue
        cv2.rectangle(O, (int(u - hw * 40 * f), int(v - hh * 40 * f)), (int(u + hw * 40 * f), int(v + hh * 40 * f)), (255, 255, 255), 1)
        cv2.putText(O, name, (int(u - 4), int(v + 5)), 0, 0.5, (255, 255, 255), 2, cv2.LINE_AA)
    cv2.putText(L, label, (3, 13), 0, 0.42, (0, 255, 255), 1, cv2.LINE_AA)
    pic = np.zeros((T, 2 * T, 3), np.uint8); pic[:L.shape[0], :L.shape[1]] = L; pic[:O.shape[0], T:T + O.shape[1]] = O
    return pic
def render(c, i):
    nodes, vias = path(c) if not c.get("control") else (c["nodes"], [])
    pa, pb = pads[c["a"]], pads[c["b"]]; core = lambda p: p.get("core") or p
    pts = [(pa["x"], pa["y"]), (pb["x"], pb["y"])] + [tuple(via_pos[v]) for v in vias if v in via_pos]
    x0, x1 = min(p[0] for p in pts) - 1.2, max(p[0] for p in pts) + 1.2; y0, y1 = min(p[1] for p in pts) - 1.2, max(p[1] for p in pts) + 1.2
    span = max(x1 - x0, y1 - y0, 4.0); cx, cy = (x0 + x1) / 2, (y0 + y1) / 2; x0, x1, y0, y1 = cx - span / 2, cx + span / 2, cy - span / 2, cy + span / 2
    id_of = lambda s: {r["id"] for r in piece_rows if r["side"] == s and r["node"] in nodes}
    marks = [("A", (core(pa)["x"], core(pa)["y"]), core(pa)["w"] / 2, core(pa)["h"] / 2), ("B", (core(pb)["x"], core(pb)["y"]), core(pb)["w"] / 2, core(pb)["h"] / 2)] + [("v", via_pos[v], 0, 0) for v in vias if v in via_pos]
    row1 = np.hstack([tile(s, x0, y0, x1, y1, id_of(s), marks, s) for s in ("top", "bottom")])
    ins = []
    for lab_, pin, p in (("A", c["a"], pa), ("B", c["b"], pb)):
        cs = sorted([x for x in contacts[pin] if x["node"] in nodes] or contacts[pin], key=lambda x: x["tier"]); s = cs[0]["side"]; m = max(core(p)["w"], core(p)["h"]) / 2 + 0.8
        ins.append(tile(s, core(p)["x"] - m, core(p)["y"] - m, core(p)["x"] + m, core(p)["y"] + m, id_of(s), [(lab_, (core(p)["x"], core(p)["y"]), core(p)["w"] / 2, core(p)["h"] / 2)], f"{lab_} {pin} {s} {cs[0]['kind']}"))
    head = np.zeros((20, 4 * T, 3), np.uint8)
    cv2.putText(head, f"#{i} {c['a']} to {c['b']}  vias on path {len(vias)}  overview span {span:.1f} mm", (4, 14), 0, 0.45, (255, 255, 255), 1, cv2.LINE_AA)
    cv2.imwrite(os.path.join(a.out, f"{i:03d}.png"), np.vstack([head, row1, np.hstack(ins)])); c["vias"] = vias; c["image"] = f"{i:03d}.png"
for i, c in enumerate(sample): render(c, i)
json.dump({"name": a.name, "population": dict(("/".join((k[0], "big" if k[1] else "small")), v) for k, v in __import__("collections").Counter((c["grade"], c["groupPins"] >= BIG) for c in conns).items()), "seed": seed, "graphSha256": hashlib.sha256(raw).hexdigest(), "populationBig": sum(1 for c in conns if c["groupPins"] >= BIG), "populationSmall": sum(1 for c in conns if c["groupPins"] < BIG), "connections": sample}, open(os.path.join(a.out, "sample.json"), "w"), indent=1)
from collections import Counter
print(a.name, "seed", seed[:16], "big", len(big), "small", len(small), "of", len(conns), Counter((c["grade"], c["groupPins"] >= BIG) for c in conns), "sample", Counter((c["grade"], c["groupPins"] >= BIG) for c in sample))