netlist.py

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 <image_path> <ocr_json>")
        print("Example: python netlist.py sample5.jpg sample5_results_ocr.json")
        sys.exit(1)

    result = generate_netlist(sys.argv[1], sys.argv[2])