import cv2 import numpy as np import json import sys import os from collections import defaultdict from datetime import datetime # ───────────────────────────────────────────── # CONFIGURATION # ───────────────────────────────────────────── PROXIMITY_THRESHOLD = 40 # pixels — tightened to reduce false connections MIN_TRACE_AREA = 150 # raised significantly to avoid false overlaps TRACE_DILATE = 2 # reduced dilation to prevent region bleeding # Reference designator prefixes per component type REFDES_MAP = { "resistor": "R", "capacitor": "C", "electrolytic capacitor": "C", "inductor": "L", "ic": "U", "transistor": "Q", "diode": "D", "led": "D", "connector": "J", "jumper": "JP", "button": "SW", "clock": "X", "transformer": "T", "potentiometer": "RV", "heatsink": "HS", "fuse": "F", "ferrite_bead": "FB", "buzzer": "BZ", "display": "DS", "battery": "BT", "emi_filter": "Z", "resistor network": "RN", "resistor jumper": "R", "capacitor jumper": "C", "pins": "TP", } # Pin counts per component type (simplified) PIN_COUNT_MAP = { "resistor": 2, "capacitor": 2, "electrolytic capacitor": 2, "inductor": 2, "diode": 2, "led": 2, "transistor": 3, "ic": 8, # default, actual varies "connector": 4, "jumper": 2, "button": 2, "clock": 4, "transformer": 4, "potentiometer": 3, "fuse": 2, "ferrite_bead": 2, "buzzer": 2, "display": 4, "battery": 2, "emi_filter": 3, "resistor network": 8, "resistor jumper": 2, "capacitor jumper": 2, "pins": 1, } # ───────────────────────────────────────────── # STEP 1 — ASSIGN REFERENCE DESIGNATORS # ───────────────────────────────────────────── def assign_reference_designators(detections: list) -> list: """ Assigns unique refdes to each component: R1, R2, C1, C2, U1, U2 etc. """ counters = defaultdict(int) annotated = [] for det in detections: label = det['label'].lower() prefix = REFDES_MAP.get(label, "X") counters[prefix] += 1 refdes = f"{prefix}{counters[prefix]}" annotated.append({ **det, "refdes": refdes, "pins": PIN_COUNT_MAP.get(label, 2), "center": get_center(det['bbox']), }) return annotated def get_center(bbox): x1, y1, x2, y2 = bbox return ((x1 + x2) // 2, (y1 + y2) // 2) # ───────────────────────────────────────────── # STEP 2 — EXTRACT TRACE MASK (OpenCV) # ───────────────────────────────────────────── def extract_trace_mask(img: np.ndarray) -> np.ndarray: """ Extracts copper traces using HSV masking + skeletonization to get thin trace lines """ hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) # Green PCB background lower_green = np.array([35, 40, 40]) upper_green = np.array([85, 255, 255]) green_mask = cv2.inRange(hsv, lower_green, upper_green) # Non-green = components + traces non_green = cv2.bitwise_not(green_mask) # Remove very dark regions (shadows/holes) lower_dark = np.array([0, 0, 0]) upper_dark = np.array([180, 255, 35]) dark_mask = cv2.inRange(hsv, lower_dark, upper_dark) trace_mask = cv2.bitwise_and(non_green, cv2.bitwise_not(dark_mask)) # Clean up kernel = np.ones((2, 2), np.uint8) trace_mask = cv2.morphologyEx(trace_mask, cv2.MORPH_OPEN, kernel, iterations=1) trace_mask = cv2.morphologyEx(trace_mask, cv2.MORPH_CLOSE, kernel, iterations=2) return trace_mask # ───────────────────────────────────────────── # STEP 3 — TRACE-BASED CONNECTION FINDING # ───────────────────────────────────────────── def find_trace_connections(components: list, trace_mask: np.ndarray, img_shape: tuple) -> list: """ For each component, dilates its bounding box region and checks if the dilated region overlaps with another component's dilated region via the trace mask. Returns list of (refdes_a, refdes_b, method) tuples. """ connections = [] n = len(components) h, w = img_shape[:2] # Build component masks comp_masks = [] for comp in components: x1, y1, x2, y2 = comp['bbox'] mask = np.zeros((h, w), dtype=np.uint8) # Dilate bbox to reach nearby traces pad = TRACE_DILATE + 5 x1p = max(0, x1 - pad) y1p = max(0, y1 - pad) x2p = min(w, x2 + pad) y2p = min(h, y2 + pad) mask[y1p:y2p, x1p:x2p] = 255 # AND with trace mask to get only trace pixels near this component comp_trace = cv2.bitwise_and(mask, trace_mask) comp_masks.append(comp_trace) # Check pairwise overlap via traces for i in range(n): for j in range(i + 1, n): # Do the trace regions of these two components overlap? overlap = cv2.bitwise_and(comp_masks[i], comp_masks[j]) if np.count_nonzero(overlap) > MIN_TRACE_AREA: connections.append(( components[i]['refdes'], components[j]['refdes'], "trace" )) return connections # ───────────────────────────────────────────── # STEP 4 — PROXIMITY-BASED CONNECTION FINDING # ───────────────────────────────────────────── def find_proximity_connections(components: list, existing_connections: list) -> list: """ Fallback: components whose centers are within PROXIMITY_THRESHOLD pixels are considered connected. Only adds connections not already found by trace method. """ existing_pairs = set( (a, b) for a, b, _ in existing_connections ) | set( (b, a) for a, b, _ in existing_connections ) new_connections = [] n = len(components) for i in range(n): for j in range(i + 1, n): a = components[i] b = components[j] pair = (a['refdes'], b['refdes']) if pair in existing_pairs or (pair[1], pair[0]) in existing_pairs: continue cx1, cy1 = a['center'] cx2, cy2 = b['center'] dist = np.sqrt((cx1 - cx2)**2 + (cy1 - cy2)**2) if dist <= PROXIMITY_THRESHOLD: new_connections.append(( a['refdes'], b['refdes'], "proximity" )) return new_connections # ───────────────────────────────────────────── # STEP 5 — BUILD NETS # ───────────────────────────────────────────── def build_nets(connections: list) -> dict: """ Groups connections into named nets using union-find style merging. Returns dict: net_name → list of refdes """ # Build adjacency adj = defaultdict(set) for a, b, method in connections: adj[a].add(b) adj[b].add(a) # Find connected components (nets) via BFS visited = set() nets = {} net_num = 1 all_nodes = set(a for a, b, _ in connections) | \ set(b for a, b, _ in connections) for node in sorted(all_nodes): if node in visited: continue # BFS queue = [node] cluster = [] while queue: curr = queue.pop(0) if curr in visited: continue visited.add(curr) cluster.append(curr) queue.extend(adj[curr] - visited) if len(cluster) > 1: net_name = f"Net-{net_num:03d}" nets[net_name] = sorted(cluster) net_num += 1 return nets # ───────────────────────────────────────────── # STEP 6 — GENERATE KICAD NETLIST (.net) # ───────────────────────────────────────────── def generate_kicad_netlist(components: list, nets: dict, output_path: str): """ Outputs a KiCAD legacy netlist .net file """ timestamp = datetime.now().strftime("%Y%m%d %H%M%S") lines = [] lines.append("(export (version D)") lines.append(f" (design") lines.append(f" (source \"pcb_image\")") lines.append(f" (date \"{timestamp}\")") lines.append(f" (tool \"PCB Image2Schematic\")") lines.append(f" )") # Components section lines.append(" (components") for comp in components: refdes = comp['refdes'] label = comp['label'] part = comp.get('part_number', 'unknown') x, y = comp['center'] lines.append(f" (comp (ref \"{refdes}\")") lines.append(f" (value \"{part}\")") lines.append(f" (description \"{label}\")") lines.append(f" (footprint \"\")") lines.append(f" (fields") lines.append(f" (field (name \"Position\") \"{x},{y}\")") lines.append(f" (field (name \"Confidence\") \"{comp['confidence']:.0%}\")") lines.append(f" )") lines.append(f" )") lines.append(" )") # Nets section lines.append(" (nets") for net_name, members in nets.items(): lines.append(f" (net (name \"{net_name}\")") for refdes in members: lines.append(f" (node (ref \"{refdes}\") (pin \"1\"))") lines.append(f" )") lines.append(" )") lines.append(")") with open(output_path, "w") as f: f.write("\n".join(lines)) print(f"[OK] KiCAD netlist saved: {output_path}") # ───────────────────────────────────────────── # PRINT NETLIST SUMMARY # ───────────────────────────────────────────── def print_netlist_summary(components: list, connections: list, nets: dict): trace_conns = [c for c in connections if c[2] == "trace"] prox_conns = [c for c in connections if c[2] == "proximity"] print(f"\n-- Netlist Summary -----------------------") print(f" Total components : {len(components)}") print(f" Total connections : {len(connections)}") print(f" via traces : {len(trace_conns)}") print(f" via proximity : {len(prox_conns)}") print(f" Total nets : {len(nets)}") print(f"\n Nets:") for net_name, members in nets.items(): print(f" {net_name}: {', '.join(members)}") print(f"------------------------------------------\n") # ───────────────────────────────────────────── # MAIN PIPELINE # ───────────────────────────────────────────── def generate_netlist(image_path: str, ocr_json_path: str) -> dict: print(f"\n{'='*50}") print(f" Netlist Generator") print(f" Image : {image_path}") print(f" Input : {ocr_json_path}") print(f"{'='*50}\n") # Load OCR results with open(ocr_json_path) as f: data = json.load(f) detections = data.get("components", []) for d in detections: d['bbox'] = tuple(d['bbox']) print(f"[OK] Loaded {len(detections)} components") # Load image for trace detection img = cv2.imread(image_path) if img is None: print(f"[X] Could not load image: {image_path}") return {} # Step 1 — assign refdes components = assign_reference_designators(detections) print(f"[OK] Reference designators assigned") for c in components: print(f" {c['refdes']:<6} — {c['label']}") # Step 2 — extract traces print(f"\n[->] Extracting trace mask...") trace_mask = extract_trace_mask(img) trace_px = np.count_nonzero(trace_mask) print(f"[OK] Trace mask extracted: {trace_px} trace pixels found") # Step 3 — trace connections print(f"\n[->] Finding trace-based connections...") trace_connections = find_trace_connections(components, trace_mask, img.shape) print(f"[OK] {len(trace_connections)} connections found via traces") # Step 4 — proximity connections (fallback) print(f"\n[->] Finding proximity-based connections...") prox_connections = find_proximity_connections(components, trace_connections) print(f"[OK] {len(prox_connections)} additional connections via proximity") all_connections = trace_connections + prox_connections # Step 5 — build nets nets = build_nets(all_connections) print(f"[OK] {len(nets)} nets built") # Step 6 — print summary print_netlist_summary(components, all_connections, nets) # Step 7 — save outputs base = os.path.splitext(image_path)[0] netlist_path = base + "_netlist.net" json_path = base + "_netlist.json" generate_kicad_netlist(components, nets, netlist_path) # Save JSON version too out_data = { "total_components": len(components), "total_connections": len(all_connections), "total_nets": len(nets), "components": [ {**c, "bbox": list(c["bbox"]), "ocr_text": [[t, conf] for t, conf in c.get("ocr_text", [])]} for c in components ], "connections": [ {"from": a, "to": b, "method": m} for a, b, m in all_connections ], "nets": { name: members for name, members in nets.items() } } with open(json_path, "w") as f: json.dump(out_data, f, indent=2) print(f"[OK] Netlist JSON saved: {json_path}") return out_data # ───────────────────────────────────────────── # ENTRY POINT # ───────────────────────────────────────────── if __name__ == "__main__": if len(sys.argv) < 3: print("Usage: python netlist.py ") print("Example: python netlist.py sample5.jpg sample5_results_ocr.json") sys.exit(1) result = generate_netlist(sys.argv[1], sys.argv[2])