Upload app.py with huggingface_hub

app.py

@@ -0,0 +1,393 @@
+import gradio as gr
+import cv2
+import numpy as np
+import json
+import os
+import tempfile
+import shutil
+from PIL import Image
+from collections import Counter
+from inference_sdk import InferenceHTTPClient
+import easyocr
+
+# ─────────────────────────────────────────────
+#  IMPORTS FROM OUR PIPELINE
+# ─────────────────────────────────────────────
+from detector import run_detection, detect_traces, draw_detections
+from ocr import run_ocr_on_detections, print_ocr_summary
+from netlist import (assign_reference_designators,
+                     extract_trace_mask,
+                     find_trace_connections,
+                     find_proximity_connections,
+                     build_nets)
+from kicad_writer import generate_kicad_schematic
+
+# ─────────────────────────────────────────────
+#  GLOBAL OCR READER (load once)
+# ─────────────────────────────────────────────
+print("[->] Loading EasyOCR...")
+ocr_reader = easyocr.Reader(['en'], gpu=False)  # CPU for HuggingFace
+print("[OK] EasyOCR ready")
+
+
+# ─────────────────────────────────────────────
+#  HELPER — numpy image to PIL
+# ─────────────────────────────────────────────
+def to_pil(img_bgr: np.ndarray) -> Image.Image:
+    return Image.fromarray(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB))
+
+
+# ─────────────────────────────────────────────
+#  TAB 1 — COMPONENT DETECTION
+# ─────────────────────────────────────────────
+def run_detection_tab(image: Image.Image):
+    if image is None:
+        return None, "❌ Please upload a PCB image first.", "{}"
+
+    # Save uploaded image to temp file
+    tmp_dir  = tempfile.mkdtemp()
+    img_path = os.path.join(tmp_dir, "input.jpg")
+    image.save(img_path)
+
+    try:
+        # Run Roboflow detection
+        detections = run_detection(img_path)
+
+        if not detections:
+            return image, "⚠️ No components detected. Try a clearer PCB image.", "{}"
+
+        # Draw detections
+        img_bgr   = cv2.imread(img_path)
+        annotated = draw_detections(img_bgr, detections)
+        result_pil= to_pil(annotated)
+
+        # Build summary text
+        counts   = Counter(d['label'] for d in detections)
+        summary  = f"✅ **{len(detections)} components detected**\n\n"
+        summary += "| Component | Count |\n|-----------|-------|\n"
+        for label, count in sorted(counts.items(), key=lambda x: -x[1]):
+            summary += f"| {label} | {count} |\n"
+
+        # Save detections to JSON for next tabs
+        det_json = json.dumps({
+            "image_path": img_path,
+            "components": [
+                {**d, "bbox": list(d["bbox"])}
+                for d in detections
+            ]
+        }, indent=2)
+
+        return result_pil, summary, det_json
+
+    except Exception as e:
+        return image, f"❌ Error: {str(e)}", "{}"
+
+
+# ─────────────────────────────────────────────
+#  TAB 2 — OCR / PART NUMBER READING
+# ─────────────────────────────────────────────
+def run_ocr_tab(det_json: str):
+    if not det_json or det_json == "{}":
+        return "⚠️ Run Detection first!", "{}"
+
+    try:
+        data       = json.loads(det_json)
+        img_path   = data.get("image_path")
+        detections = data.get("components", [])
+
+        for d in detections:
+            d['bbox'] = tuple(d['bbox'])
+
+        IC_LABELS = ['ic', 'transistor', 'clock', 'display']
+
+        # Run OCR using global reader
+        img = cv2.imread(img_path)
+        ih, iw = img.shape[:2]
+        PADDING = 10
+        updated = []
+
+        for det in detections:
+            label = det['label']
+            if label not in IC_LABELS:
+                det['ocr_text']    = []
+                det['part_number'] = "N/A"
+                updated.append(det)
+                continue
+
+            x1, y1, x2, y2 = det['bbox']
+            x1p = max(0,  x1 - PADDING)
+            y1p = max(0,  y1 - PADDING)
+            x2p = min(iw, x2 + PADDING)
+            y2p = min(ih, y2 + PADDING)
+            patch = img[y1p:y2p, x1p:x2p]
+
+            if patch.size == 0:
+                det['ocr_text']    = []
+                det['part_number'] = "unknown"
+                updated.append(det)
+                continue
+
+            # Upscale for better OCR
+            h, w   = patch.shape[:2]
+            scale  = 3 if max(h, w) < 100 else 2
+            patch  = cv2.resize(patch, (w*scale, h*scale),
+                                interpolation=cv2.INTER_CUBIC)
+
+            results  = ocr_reader.readtext(patch)
+            texts    = [(t.strip(), round(c, 3))
+                        for _, t, c in results
+                        if c >= 0.4 and len(t.strip()) >= 2]
+
+            combined           = " ".join(t for t, c in texts).strip()
+            det['ocr_text']    = texts
+            det['part_number'] = combined if combined else "unknown"
+            updated.append(det)
+
+        # Build output table
+        ic_dets    = [d for d in updated if d['label'] in IC_LABELS]
+        identified = [d for d in ic_dets
+                      if d.get('part_number', 'unknown') not in
+                      ('unknown', 'N/A', '')]
+
+        summary  = f"✅ **OCR complete — {len(identified)}/{len(ic_dets)} ICs identified**\n\n"
+        summary += "| RefDes | Label | Part Number | Confidence |\n"
+        summary += "|--------|-------|-------------|------------|\n"
+
+        for i, det in enumerate(ic_dets):
+            ref  = f"U{i+1}"
+            part = det.get('part_number', 'unknown')
+            conf = det['confidence']
+            summary += f"| {ref} | {det['label']} | {part} | {conf:.0%} |\n"
+
+        # Pass updated detections forward
+        out_json = json.dumps({
+            "image_path": img_path,
+            "components": [
+                {**d,
+                 "bbox":     list(d["bbox"]),
+                 "ocr_text": [[t, c] for t, c in d.get("ocr_text", [])]}
+                for d in updated
+            ]
+        }, indent=2)
+
+        return summary, out_json
+
+    except Exception as e:
+        return f"❌ Error: {str(e)}", "{}"
+
+
+# ─────────────────────────────────────────────
+#  TAB 3 — NETLIST GENERATION
+# ─────────────────────────────────────────────
+def run_netlist_tab(ocr_json: str):
+    if not ocr_json or ocr_json == "{}":
+        return "⚠️ Run OCR first!", "{}", None
+
+    try:
+        data       = json.loads(ocr_json)
+        img_path   = data.get("image_path")
+        detections = data.get("components", [])
+
+        for d in detections:
+            d['bbox'] = tuple(d['bbox'])
+
+        img        = cv2.imread(img_path)
+        components = assign_reference_designators(detections)
+        trace_mask = extract_trace_mask(img)
+        trace_conn = find_trace_connections(components, trace_mask, img.shape)
+        prox_conn  = find_proximity_connections(components, trace_conn)
+        all_conn   = trace_conn + prox_conn
+        nets       = build_nets(all_conn)
+
+        # Save netlist JSON to temp file
+        tmp_dir      = os.path.dirname(img_path)
+        netlist_path = os.path.join(tmp_dir, "netlist.json")
+
+        out_data = {
+            "total_components": len(components),
+            "total_connections": len(all_conn),
+            "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_conn
+            ],
+            "nets": nets
+        }
+
+        with open(netlist_path, "w") as f:
+            json.dump(out_data, f, indent=2)
+
+        # Build summary
+        trace_c = len([c for c in all_conn if c[2] == "trace"])
+        prox_c  = len([c for c in all_conn if c[2] == "proximity"])
+
+        summary  = f"✅ **Netlist generated successfully**\n\n"
+        summary += f"- **Components:** {len(components)}\n"
+        summary += f"- **Connections:** {len(all_conn)} "
+        summary += f"({trace_c} via traces, {prox_c} via proximity)\n"
+        summary += f"- **Nets:** {len(nets)}\n\n"
+        summary += "| Net | Members |\n|-----|--------|\n"
+        for net_name, members in nets.items():
+            summary += f"| {net_name} | {', '.join(members[:5])}"
+            if len(members) > 5:
+                summary += f" ... +{len(members)-5} more"
+            summary += " |\n"
+
+        return summary, json.dumps({"netlist_path": netlist_path}), netlist_path
+
+    except Exception as e:
+        return f"❌ Error: {str(e)}", "{}", None
+
+
+# ────────────────────────��────────────────────
+#  TAB 4 — KICAD SCHEMATIC OUTPUT
+# ─────────────────────────────────────────────
+def run_kicad_tab(netlist_ref: str):
+    if not netlist_ref or netlist_ref == "{}":
+        return "⚠️ Run Netlist generation first!", None
+
+    try:
+        data         = json.loads(netlist_ref)
+        netlist_path = data.get("netlist_path")
+
+        if not netlist_path or not os.path.exists(netlist_path):
+            return "❌ Netlist file not found. Re-run previous steps.", None
+
+        # Generate KiCAD schematic
+        tmp_dir    = os.path.dirname(netlist_path)
+        sch_path   = os.path.join(tmp_dir, "schematic.kicad_sch")
+
+        generate_kicad_schematic(netlist_path, sch_path)
+
+        summary  = f"✅ **KiCAD schematic generated!**\n\n"
+        summary += f"- File: `schematic.kicad_sch`\n"
+        summary += f"- Format: KiCAD 6/7 compatible\n\n"
+        summary += "**How to open:**\n"
+        summary += "1. Download the file below\n"
+        summary += "2. Open KiCAD → File → Open Schematic\n"
+        summary += "   OR drag into [kicanvas.org](https://kicanvas.org) for instant preview\n"
+
+        return summary, sch_path
+
+    except Exception as e:
+        return f"❌ Error: {str(e)}", None
+
+
+# ─────────────────────────────────────────────
+#  BUILD GRADIO UI
+# ─────────────────────────────────────────────
+def build_ui():
+    with gr.Blocks(
+        title="PCB Image → Schematic",
+        theme=gr.themes.Soft(),
+        css="""
+        .tab-header { font-size: 1.1em; font-weight: bold; }
+        .output-panel { background: #1a1a2e; border-radius: 8px; }
+        """
+    ) as demo:
+
+        # ── Header ──
+        gr.Markdown("""
+        # 🔌 PCB Image → Schematic
+        ### Convert a PCB photo into a KiCAD schematic automatically
+        Upload a PCB image and step through each stage of the pipeline.
+        """)
+
+        # ── Shared state between tabs ──
+        detection_state = gr.State("{}")
+        ocr_state       = gr.State("{}")
+        netlist_state   = gr.State("{}")
+
+        # ── Tab 1: Detection ──
+        with gr.Tab("📷 1 — Component Detection"):
+            gr.Markdown("Upload a PCB image or click one of the example images below.")
+            with gr.Row():
+                with gr.Column(scale=1):
+                    img_input  = gr.Image(type="pil", label="PCB Image")
+                    detect_btn = gr.Button("🔍 Detect Components", variant="primary")
+                    gr.Examples(
+                        examples=[
+                            ["sample 1.jpg"],
+                            ["sample 2.jpg"],
+                            ["sample 3.jpg"],
+                            ["sample 4.jpg"],
+                            ["sample 5.jpg"],
+                        ],
+                        inputs=img_input,
+                        label="📂 Example PCB Images — click to load",
+                        examples_per_page=5,
+                    )
+                with gr.Column(scale=1):
+                    detect_out  = gr.Image(label="Detected Components")
+                    detect_text = gr.Markdown()
+
+            detect_btn.click(
+                fn=run_detection_tab,
+                inputs=[img_input],
+                outputs=[detect_out, detect_text, detection_state]
+            )
+
+        # ── Tab 2: OCR ──
+        with gr.Tab("🔤 2 — Read IC Text"):
+            gr.Markdown("Reads part numbers from IC chips using OCR.")
+            ocr_btn  = gr.Button("📖 Run OCR on ICs", variant="primary")
+            ocr_text = gr.Markdown()
+
+            ocr_btn.click(
+                fn=run_ocr_tab,
+                inputs=[detection_state],
+                outputs=[ocr_text, ocr_state]
+            )
+
+        # ── Tab 3: Netlist ──
+        with gr.Tab("🔗 3 — Generate Netlist"):
+            gr.Markdown("Finds connections between components using trace detection + proximity.")
+            netlist_btn  = gr.Button("⚡ Generate Netlist", variant="primary")
+            netlist_text = gr.Markdown()
+            netlist_file = gr.File(label="Download Netlist JSON", visible=False)
+
+            netlist_btn.click(
+                fn=run_netlist_tab,
+                inputs=[ocr_state],
+                outputs=[netlist_text, netlist_state, netlist_file]
+            )
+
+        # ── Tab 4: KiCAD ──
+        with gr.Tab("📐 4 — KiCAD Schematic"):
+            gr.Markdown("Generates a KiCAD `.kicad_sch` file you can open in KiCAD or kicanvas.org")
+            kicad_btn  = gr.Button("💾 Generate KiCAD File", variant="primary")
+            kicad_text = gr.Markdown()
+            kicad_file = gr.File(label="Download .kicad_sch")
+
+            kicad_btn.click(
+                fn=run_kicad_tab,
+                inputs=[netlist_state],
+                outputs=[kicad_text, kicad_file]
+            )
+
+        # ── Footer ──
+        gr.Markdown("""
+        ---
+        Built with Roboflow YOLOv8 · EasyOCR · OpenCV · KiCAD
+        """)
+
+    return demo
+
+
+# ─────────────────────────────────────────────
+#  ENTRY POINT
+# ─────────────────────────────────────────────
+if __name__ == "__main__":
+    demo = build_ui()
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    )