Spaces:
Sleeping
Sleeping
File size: 5,128 Bytes
add24bd 5673379 f63e98a 5673379 add24bd 5673379 f63e98a add24bd f63e98a add24bd f63e98a 5673379 add24bd 4dfed21 add24bd 5673379 add24bd f63e98a add24bd f63e98a add24bd 5673379 add24bd 5673379 add24bd 5673379 4dfed21 5673379 add24bd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | """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,
)
|