"""PCB defect detection - Hugging Face Space (Gradio). Self-contained: no torch/ultralytics/pcb_defect import, only onnxruntime + cv2 + gradio + huggingface_hub. The letterbox/preprocess/postprocess functions below are a deliberate near-duplicate (~60 lines) of src/pcb_defect/e2e_onnx.py's logic - the Space can't easily depend on the full project package, so this file stands alone. scripts/verify_onnx_parity.py's parity gate covers this same logic (see plan.md SS 2.3/2.6): any change here should be mirrored there and re-verified. Env vars: MODEL_REPO HF model repo id to pull best.onnx from (default below). MODEL_PATH_OVERRIDE local file path - skips the HF download, for testing this Space before MODEL_REPO actually has anything uploaded to it (see plan.md SS 2.6/2.7 ordering). """ from __future__ import annotations import os import time from dataclasses import dataclass import cv2 import gradio as gr import numpy as np from PIL import Image MODEL_REPO = os.environ.get("MODEL_REPO", "betty0/pcb-defect-detection") MODEL_PATH_OVERRIDE = os.environ.get("MODEL_PATH_OVERRIDE") CLASSES = ["missing_hole", "mouse_bite", "open_circuit", "short", "spur", "spurious_copper"] IMG_SIZE = 640 PAD_VALUE = 114 DEFAULT_CONF = 0.25 EXAMPLES_DIR = os.path.join(os.path.dirname(__file__), "examples") @dataclass class LetterboxInfo: gain: float pad_left: float pad_top: float @dataclass class Detection: cls_id: int xyxy: tuple[float, float, float, float] conf: float def letterbox(image: Image.Image, size: int = IMG_SIZE) -> tuple[np.ndarray, LetterboxInfo]: """Matches ultralytics' LetterBox (cv2.resize + cv2.copyMakeBorder) exactly - PIL's resize is NOT numerically interchangeable with cv2's (verified empirically, see plan.md SS 2.3).""" rgb = np.asarray(image.convert("RGB")) h, w = rgb.shape[:2] gain = min(size / w, size / h) new_w, new_h = round(w * gain), round(h * gain) resized = cv2.resize(rgb, (new_w, new_h), interpolation=cv2.INTER_LINEAR) dw, dh = (size - new_w) / 2, (size - new_h) / 2 top, bottom = round(dh - 0.1), round(dh + 0.1) left, right = round(dw - 0.1), round(dw + 0.1) canvas = cv2.copyMakeBorder( resized, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(PAD_VALUE,) * 3 ) return canvas, LetterboxInfo(gain, left, top) def preprocess(image: Image.Image) -> tuple[np.ndarray, LetterboxInfo]: canvas, info = letterbox(image) chw = canvas.transpose(2, 0, 1).astype(np.float32) / 255.0 batch = np.ascontiguousarray(np.expand_dims(chw, axis=0)) return batch, info def postprocess( output: np.ndarray, info: LetterboxInfo, orig_size: tuple[int, int], conf: float ) -> list[Detection]: """(1, 300, 6) letterboxed-space rows -> Detection list in original-image pixel coords.""" rows = output[0] rows = rows[rows[:, 4] >= conf] orig_w, orig_h = orig_size detections = [] for x1, y1, x2, y2, score, cls in rows: ox1 = float(max(0.0, min((x1 - info.pad_left) / info.gain, orig_w))) oy1 = float(max(0.0, min((y1 - info.pad_top) / info.gain, orig_h))) ox2 = float(max(0.0, min((x2 - info.pad_left) / info.gain, orig_w))) oy2 = float(max(0.0, min((y2 - info.pad_top) / info.gain, orig_h))) detections.append(Detection(int(cls), (ox1, oy1, ox2, oy2), float(score))) return detections def _load_session(): import onnxruntime as ort if MODEL_PATH_OVERRIDE: model_path = MODEL_PATH_OVERRIDE else: from huggingface_hub import hf_hub_download model_path = hf_hub_download(repo_id=MODEL_REPO, filename="best.onnx") session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"]) return session, session.get_inputs()[0].name _SESSION, _INPUT_NAME = _load_session() def _to_annotations(detections: list[Detection]) -> list[tuple[tuple[int, int, int, int], str]]: return [ ( (round(d.xyxy[0]), round(d.xyxy[1]), round(d.xyxy[2]), round(d.xyxy[3])), f"{CLASSES[d.cls_id]} {d.conf:.2f}", ) for d in detections ] def _to_table(detections: list[Detection]) -> list[list]: return [ [CLASSES[d.cls_id], round(d.conf, 3), *[round(c, 1) for c in d.xyxy]] for d in detections ] def run_inference(image: Image.Image | None, conf: float): if image is None: return None, None, "", [] t0 = time.perf_counter() batch, info = preprocess(image) (raw_output,) = _SESSION.run(None, {_INPUT_NAME: batch}) elapsed_ms = (time.perf_counter() - t0) * 1000 cache = {"raw_output": raw_output, "info": info, "orig_size": image.size, "image": image} detections = postprocess(raw_output, info, image.size, conf) annotated = (image, _to_annotations(detections)) latency_text = f"推論時間：{elapsed_ms:.0f} ms（僅第一次上傳/換圖需要；拖曳滑桿不重新推論）" return annotated, cache, latency_text, _to_table(detections) def rerender(cache: dict | None, conf: float): if cache is None: return None, [] detections = postprocess(cache["raw_output"], cache["info"], cache["orig_size"], conf) return (cache["image"], _to_annotations(detections)), _to_table(detections) with gr.Blocks(title="PCB 裸板瑕疵偵測 - YOLO26 Demo") as demo: gr.Markdown( "# PCB 裸板瑕疵偵測（YOLO26，NMS-free e2e，ONNX Runtime CPU）\n\n" "上傳一張 PCB 裸板照片，或點選下方範例。拖曳信心值滑桿只會重新篩選、" "不會重新跑推論（後端快取了原始 (300,6) 輸出）。" f"6 類瑕疵：{'、'.join(CLASSES)}。" ) with gr.Row(): with gr.Column(): image_input = gr.Image(type="pil", label="上傳 PCB 影像") conf_slider = gr.Slider( minimum=0.05, maximum=0.90, value=DEFAULT_CONF, step=0.01, label="信心值門檻" ) with gr.Column(): annotated_output = gr.AnnotatedImage(label="偵測結果") latency_text = gr.Markdown() results_table = gr.Dataframe( headers=["類別", "信心值", "x1", "y1", "x2", "y2"], label="偵測列表", interactive=False, ) raw_state = gr.State(value=None) image_input.upload( fn=run_inference, inputs=[image_input, conf_slider], outputs=[annotated_output, raw_state, latency_text, results_table], ) # gr.Examples only sets image_input's value - it does not fire .upload() (that event is # specifically for actual file uploads), so without run_on_click=True, clicking an example # would show no detections until some other action ran inference (confirmed against # gradio/helpers.py: cache_examples=False + run_on_click=False silently skips fn entirely). # run_on_click=True runs inference live on each click - explicit rather than relying on # HF Spaces' implicit cache_examples=True default, so behavior matches local testing too. gr.Examples( examples=[[os.path.join(EXAMPLES_DIR, f"04_{cls}_01.jpg")] for cls in CLASSES], inputs=[image_input], outputs=[annotated_output, raw_state, latency_text, results_table], fn=lambda img: run_inference(img, DEFAULT_CONF), run_on_click=True, cache_examples=False, label="範例（每類一張）", ) conf_slider.release( fn=rerender, inputs=[raw_state, conf_slider], outputs=[annotated_output, results_table], ) if __name__ == "__main__": demo.launch()