Spaces:
Sleeping
Sleeping
| """Gradio demo: PCB defect + golden reference -> enriched image (A11_CA prebackbone).""" | |
| from __future__ import annotations | |
| import os | |
| from pathlib import Path | |
| import gradio as gr | |
| import numpy as np | |
| from PIL import Image | |
| from prebackbone_infer import enrich_pair, get_enricher | |
| EXAMPLES_DIR = Path(__file__).resolve().parent / "examples" | |
| TITLE = "RefDiffNet — PCB Reference–Defect Enrichment" | |
| DESCRIPTION = """ | |
| Upload a **defect PCB image** and its **golden reference** (same board, no defect). | |
| The trained A11_CA prebackbone fuses them into an **enriched** image used by YOLO detection. | |
| Formula: `enriched = defect + α · gate · delta` | |
| """ | |
| def _examples() -> list[list[str]]: | |
| if not EXAMPLES_DIR.exists(): | |
| return [] | |
| inputs = sorted(EXAMPLES_DIR.glob("*_input.jpg")) | |
| out: list[list[str]] = [] | |
| for inp in inputs: | |
| ref = EXAMPLES_DIR / inp.name.replace("_input.", "_reference.") | |
| if ref.exists(): | |
| out.append([str(inp), str(ref)]) | |
| return out | |
| def run_demo(defect_img, reference_img): | |
| if defect_img is None or reference_img is None: | |
| raise gr.Error("Please upload both defect and reference images.") | |
| defect_rgb, ref_rgb, enriched_rgb = enrich_pair(defect_img, reference_img) | |
| # Side-by-side comparison for the demo panel | |
| h = max(defect_rgb.shape[0], ref_rgb.shape[0], enriched_rgb.shape[0]) | |
| def _pad(im: np.ndarray) -> np.ndarray: | |
| if im.shape[0] == h: | |
| return im | |
| pad = h - im.shape[0] | |
| return np.pad(im, ((0, pad), (0, 0), (0, 0)), mode="constant", constant_values=114) | |
| row = np.concatenate([_pad(defect_rgb), _pad(ref_rgb), _pad(enriched_rgb)], axis=1) | |
| return enriched_rgb, row | |
| def _preload_model(): | |
| try: | |
| get_enricher() | |
| except FileNotFoundError as e: | |
| print(f"[warn] Model not loaded yet: {e}") | |
| if __name__ == "__main__": | |
| _preload_model() | |
| with gr.Blocks(title=TITLE) as demo: | |
| gr.Markdown(f"# {TITLE}") | |
| gr.Markdown(DESCRIPTION) | |
| with gr.Row(): | |
| defect_in = gr.Image(label="Defect image (input)", type="numpy", image_mode="RGB") | |
| ref_in = gr.Image(label="Golden reference", type="numpy", image_mode="RGB") | |
| run_btn = gr.Button("Generate enriched image", variant="primary") | |
| with gr.Row(): | |
| enriched_out = gr.Image(label="Enriched output", type="numpy", image_mode="RGB") | |
| compare_out = gr.Image( | |
| label="Defect | Reference | Enriched", | |
| type="numpy", | |
| image_mode="RGB", | |
| ) | |
| run_btn.click(fn=run_demo, inputs=[defect_in, ref_in], outputs=[enriched_out, compare_out]) | |
| ex = _examples() | |
| if ex: | |
| gr.Examples(examples=ex, inputs=[defect_in, ref_in], label="Example pairs") | |
| demo.launch( | |
| server_name="0.0.0.0", | |
| server_port=int(os.environ.get("PORT", "7860")), | |
| share=os.environ.get("GRADIO_SHARE", "").lower() in {"1", "true", "yes"}, | |
| ) | |