| """Hugging Face Space: upload a PCB image, get defect boxes from the INT8 ONNX model. |
| |
| Pure ONNX Runtime + OpenCV (no PyTorch / Ultralytics), so the Space is small and |
| runs on any hardware tier's CPU. Put the quantized model next to this file as |
| model_int8.onnx (or set MODEL_PATH). Optional: a few test-split images in examples/. |
| """ |
| import ast |
| import os |
| import time |
| from pathlib import Path |
|
|
| import cv2 |
| import gradio as gr |
| import numpy as np |
| import onnxruntime as ort |
|
|
| try: |
| import spaces |
|
|
| @spaces.GPU |
| def _unused_gpu_placeholder(): |
| """Never called: inference below runs on CPU only.""" |
| return None |
| except ImportError: |
| pass |
|
|
| HERE = Path(__file__).parent |
| MODEL_PATH = os.environ.get("MODEL_PATH", str(HERE / "model_int8.onnx")) |
| IMGSZ = int(os.environ.get("IMGSZ", "1280")) |
| NMS_IOU = 0.7 |
| FALLBACK_NAMES = ["missing_hole", "mouse_bite", "open_circuit", "short", "spur", "spurious_copper"] |
| COLORS = [(56, 56, 255), (151, 157, 255), (31, 112, 255), (29, 178, 255), (49, 210, 207), (10, 249, 72)] |
|
|
| opts = ort.SessionOptions() |
| opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL |
| session = ort.InferenceSession(MODEL_PATH, opts, providers=["CPUExecutionProvider"]) |
| INPUT_NAME = session.get_inputs()[0].name |
| _meta = session.get_modelmeta().custom_metadata_map |
| NAMES = ast.literal_eval(_meta["names"]) if "names" in _meta else dict(enumerate(FALLBACK_NAMES)) |
|
|
|
|
| def letterbox(img, size): |
| """Resize keeping aspect ratio and pad to size x size (gray 114), like Ultralytics.""" |
| h, w = img.shape[:2] |
| r = min(size / h, size / w) |
| nh, nw = round(h * r), round(w * r) |
| top, left = (size - nh) // 2, (size - nw) // 2 |
| out = np.full((size, size, 3), 114, dtype=np.uint8) |
| out[top:top + nh, left:left + nw] = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_LINEAR) |
| return out, r, left, top |
|
|
|
|
| def predict(bgr, conf): |
| """Return a list of (class_id, score, x1, y1, x2, y2) in original-image pixels.""" |
| padded, r, dx, dy = letterbox(bgr, IMGSZ) |
| x = np.ascontiguousarray(padded[:, :, ::-1].transpose(2, 0, 1), dtype=np.float32)[None] / 255.0 |
| out = session.run(None, {INPUT_NAME: x})[0][0].T |
| scores = out[:, 4:] |
| cls = scores.argmax(1) |
| score = scores[np.arange(len(cls)), cls] |
| keep = score >= conf |
| if not keep.any(): |
| return [] |
| xywh, cls, score = out[keep, :4], cls[keep], score[keep] |
| h, w = bgr.shape[:2] |
| x1 = np.clip((xywh[:, 0] - xywh[:, 2] / 2 - dx) / r, 0, w) |
| y1 = np.clip((xywh[:, 1] - xywh[:, 3] / 2 - dy) / r, 0, h) |
| x2 = np.clip((xywh[:, 0] + xywh[:, 2] / 2 - dx) / r, 0, w) |
| y2 = np.clip((xywh[:, 1] + xywh[:, 3] / 2 - dy) / r, 0, h) |
| boxes = np.stack([x1, y1, x2 - x1, y2 - y1], 1) |
| idx = cv2.dnn.NMSBoxesBatched(boxes.tolist(), score.tolist(), cls.tolist(), conf, NMS_IOU) |
| idx = np.asarray(idx, dtype=int).reshape(-1) |
| return [(int(cls[i]), float(score[i]), float(x1[i]), float(y1[i]), float(x2[i]), float(y2[i])) for i in idx] |
|
|
|
|
| def draw(bgr, dets): |
| out = bgr.copy() |
| thick = max(2, round(max(out.shape[:2]) / 600)) |
| for c, s, x1, y1, x2, y2 in dets: |
| color = COLORS[c % len(COLORS)] |
| cv2.rectangle(out, (int(x1), int(y1)), (int(x2), int(y2)), color, thick) |
| cv2.putText(out, f"{NAMES[c]} {s:.2f}", (int(x1), max(20, int(y1) - 8)), |
| cv2.FONT_HERSHEY_SIMPLEX, thick / 2.5, color, thick, cv2.LINE_AA) |
| return out |
|
|
|
|
| def detect(image, conf): |
| if image is None: |
| return None, [], "" |
| bgr = np.ascontiguousarray(image[:, :, ::-1]) |
| t0 = time.perf_counter() |
| dets = predict(bgr, conf) |
| ms = (time.perf_counter() - t0) * 1000 |
| rows = [[NAMES[c], round(s, 3), int(x1), int(y1), int(x2), int(y2)] for c, s, x1, y1, x2, y2 in dets] |
| status = f"{len(rows)} defects found in {ms:.0f} ms (CPU, ONNX Runtime, preprocessing + model + NMS)" |
| return draw(bgr, dets)[:, :, ::-1], rows, status |
|
|
|
|
| |
| EXAMPLES = {p.stem: str(p) for p in sorted([*HERE.glob("examples/*.jpg"), *HERE.glob("*.jpg")])} |
|
|
|
|
| def run_example(name, conf): |
| """Load a bundled test image and detect on it in one step.""" |
| if not name: |
| return None, None, [], "" |
| rgb = np.ascontiguousarray(cv2.imread(EXAMPLES[name])[:, :, ::-1]) |
| annotated, rows, status = detect(rgb, conf) |
| return rgb, annotated, rows, status |
|
|
| with gr.Blocks(title="PCB defect detection (INT8, CPU)") as demo: |
| gr.Markdown( |
| "# PCB defect detection\n" |
| "YOLO11n quantized to INT8 with ONNX Runtime (detection head kept in FP32), running on CPU. " |
| "Trained on PKU-Market-PCB, a public research dataset, not production factory data. " |
| "Defects are tiny: zoom into the result to see the boxes." |
| ) |
| with gr.Row(): |
| inp = gr.Image(type="numpy", label="PCB image", format="jpeg") |
| out = gr.Image(type="numpy", label="Detections", format="jpeg") |
| conf = gr.Slider(0.05, 0.9, value=0.25, step=0.05, label="Confidence threshold") |
| status = gr.Markdown() |
| table = gr.Dataframe(headers=["class", "conf", "x1", "y1", "x2", "y2"], label="Boxes") |
| if EXAMPLES: |
| pick = gr.Radio(list(EXAMPLES), value=None, label="Or pick a test-split example, then press Detect") |
| gr.Button("Detect", variant="primary").click(detect, [inp, conf], [out, table, status]) |
| if EXAMPLES: |
| pick.change(run_example, [pick, conf], [inp, out, table, status]) |
|
|
| if __name__ == "__main__": |
| demo.launch() |
|
|
