Spaces:
Sleeping
Sleeping
| """Gradio demo: PCB defect + golden reference -> enriched image (A11_CA prebackbone).""" | |
| from __future__ import annotations | |
| import os | |
| import socket | |
| from pathlib import Path | |
| _ROOT = Path(__file__).resolve().parent | |
| def _is_hf_space() -> bool: | |
| return bool(os.environ.get("SPACE_ID") or os.environ.get("SYSTEM") == "spaces") | |
| def _configure_runtime() -> None: | |
| os.environ.setdefault("GRADIO_SERVER_PORT", os.environ.get("PORT", "7860")) | |
| os.environ.setdefault( | |
| "PREBACKBONE_ONLY_WEIGHTS", | |
| str(_ROOT / "weights" / "prebackbone_a11_ca.pt"), | |
| ) | |
| if _is_hf_space(): | |
| os.environ.setdefault("PRELOAD_MODEL", "1") | |
| os.environ.setdefault("GRADIO_SERVER_NAME", "0.0.0.0") | |
| _configure_runtime() | |
| import gradio_patch | |
| gradio_patch.apply() | |
| import gradio as gr | |
| import numpy as np | |
| from PIL import Image | |
| from prebackbone_infer import enrich_pair, get_enricher | |
| TITLE = "RefDiffNet — PCB Reference–Defect Enrichment" | |
| DESCRIPTION = """ | |
| Upload a **defect PCB image** and its **golden reference** (same H×W, pre-aligned — e.g. `*_input.jpg` / `*_reference.jpg` from training). | |
| The A11_CA prebackbone outputs an **enriched** image: `enriched = defect + α · gate · delta` | |
| """ | |
| def run_demo(defect_path: str | None, reference_path: str | None): | |
| if not defect_path or not reference_path: | |
| raise gr.Error("Please upload both defect and reference images.") | |
| defect_rgb, ref_rgb, enriched_rgb = enrich_pair(defect_path, reference_path) | |
| defect_rgb = np.array(defect_rgb, dtype=np.uint8, copy=True) | |
| ref_rgb = np.array(ref_rgb, dtype=np.uint8, copy=True) | |
| enriched_rgb = np.array(enriched_rgb, dtype=np.uint8, copy=True) | |
| 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 Image.fromarray(enriched_rgb), Image.fromarray(row) | |
| def _example_pairs() -> list[list[str]]: | |
| pairs: list[list[str]] = [] | |
| ex_dir = _ROOT / "examples" | |
| if not ex_dir.is_dir(): | |
| return pairs | |
| for inp in sorted(ex_dir.glob("*_input.jpg")): | |
| ref = inp.with_name(inp.name.replace("_input.", "_reference.")) | |
| if ref.is_file(): | |
| pairs.append([str(inp), str(ref)]) | |
| return pairs | |
| def _preload_model() -> None: | |
| try: | |
| get_enricher() | |
| print("[RefDiffNet] Prebackbone loaded.") | |
| except FileNotFoundError as e: | |
| print(f"[RefDiffNet] WARN: {e}") | |
| def build_demo() -> gr.Blocks: | |
| with gr.Blocks(title=TITLE) as demo: | |
| gr.Markdown(DESCRIPTION.strip()) | |
| with gr.Row(): | |
| defect_in = gr.Image(type="filepath", label="Defect image (input)") | |
| ref_in = gr.Image(type="filepath", label="Golden reference") | |
| run_btn = gr.Button("Run enrichment", variant="primary") | |
| with gr.Row(): | |
| out_enriched = gr.Image(type="pil", label="Enriched output") | |
| out_row = gr.Image(type="pil", label="Defect | Reference | Enriched") | |
| run_btn.click(run_demo, inputs=[defect_in, ref_in], outputs=[out_enriched, out_row]) | |
| examples = _example_pairs() | |
| if examples: | |
| gr.Examples( | |
| examples=examples, | |
| inputs=[defect_in, ref_in], | |
| examples_per_page=3, | |
| label="Example pairs (pre-aligned HR-IPCB samples)", | |
| ) | |
| return demo | |
| demo = build_demo() | |
| def _launch_kwargs() -> dict: | |
| if _is_hf_space(): | |
| return { | |
| "server_name": "0.0.0.0", | |
| "server_port": int(os.environ.get("PORT", "7860")), | |
| "share": False, | |
| } | |
| server_name = os.environ.get("GRADIO_SERVER_NAME", "127.0.0.1") | |
| preferred_port = int(os.environ.get("PORT", os.environ.get("GRADIO_SERVER_PORT", "7860"))) | |
| bind_host = "127.0.0.1" if server_name in ("127.0.0.1", "localhost") else server_name | |
| server_port = preferred_port | |
| if server_name in ("127.0.0.1", "localhost"): | |
| for port in range(preferred_port, preferred_port + 20): | |
| with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: | |
| sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) | |
| try: | |
| sock.bind((bind_host, port)) | |
| server_port = port | |
| break | |
| except OSError: | |
| continue | |
| share = os.environ.get("GRADIO_SHARE", "").lower() in ("1", "true", "yes") | |
| return { | |
| "server_name": server_name, | |
| "server_port": server_port, | |
| "share": share, | |
| } | |
| if __name__ == "__main__": | |
| print(f"[RefDiffNet] gradio {getattr(gr, '__version__', 'unknown')}") | |
| if os.environ.get("PRELOAD_MODEL", "0").lower() in ("1", "true", "yes"): | |
| _preload_model() | |
| demo.launch( | |
| **(_launch_kwargs()), | |
| show_error=True, | |
| inbrowser=False, | |
| prevent_thread_lock=False, | |
| ) | |