Spaces:
Sleeping
Sleeping
| """Shared ONNX Runtime inference for the PCB defect detector. | |
| Used by both the FastAPI service (serve/) and the Gradio demo (app/app.py), so the | |
| two always run the same preprocessing, decoding and NMS. The steps mirror | |
| Ultralytics 8.3 (LetterBox, non_max_suppression, scale_boxes) so results match the | |
| evaluation pipeline in src/. | |
| Only numpy, OpenCV and onnxruntime are needed: no PyTorch, no Ultralytics. | |
| """ | |
| from __future__ import annotations | |
| import ast | |
| import time | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| import cv2 | |
| import numpy as np | |
| import onnxruntime as ort | |
| 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)] # BGR | |
| MAX_WH = 7680 # class offset for batched NMS, as in Ultralytics | |
| class Detection: | |
| class_id: int | |
| class_name: str | |
| confidence: float | |
| box: tuple[float, float, float, float] # x1, y1, x2, y2 in original-image pixels | |
| def to_dict(self) -> dict: | |
| return { | |
| "class": self.class_name, | |
| "class_id": self.class_id, | |
| "confidence": round(self.confidence, 4), | |
| "box": [round(v, 1) for v in self.box], | |
| } | |
| class Timings: | |
| preprocess_ms: float | |
| model_ms: float | |
| postprocess_ms: float | |
| def total_ms(self) -> float: | |
| return self.preprocess_ms + self.model_ms + self.postprocess_ms | |
| def letterbox(bgr: np.ndarray, size: int) -> tuple[np.ndarray, float, tuple[float, float]]: | |
| """Ultralytics LetterBox(auto=False, center=True): resize keeping aspect ratio, pad with 114.""" | |
| h, w = bgr.shape[:2] | |
| r = min(size / h, size / w) | |
| new_w, new_h = int(round(w * r)), int(round(h * r)) | |
| dw, dh = (size - new_w) / 2, (size - new_h) / 2 | |
| img = bgr if (w, h) == (new_w, new_h) else cv2.resize(bgr, (new_w, new_h), interpolation=cv2.INTER_LINEAR) | |
| top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) | |
| left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) | |
| img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114)) | |
| return img, r, (dw, dh) | |
| def nms(boxes: np.ndarray, scores: np.ndarray, iou_thres: float) -> np.ndarray: | |
| """Greedy NMS (same rule as torchvision.ops.nms). boxes: (n, 4) xyxy. Returns kept indices.""" | |
| order = scores.argsort()[::-1] | |
| x1, y1, x2, y2 = boxes.T | |
| areas = (x2 - x1) * (y2 - y1) | |
| keep = [] | |
| while order.size: | |
| i = order[0] | |
| keep.append(i) | |
| rest = order[1:] | |
| iw = np.clip(np.minimum(x2[i], x2[rest]) - np.maximum(x1[i], x1[rest]), 0, None) | |
| ih = np.clip(np.minimum(y2[i], y2[rest]) - np.maximum(y1[i], y1[rest]), 0, None) | |
| inter = iw * ih | |
| iou = inter / (areas[i] + areas[rest] - inter + 1e-9) | |
| order = rest[iou <= iou_thres] | |
| return np.asarray(keep, dtype=int) | |
| class Detector: | |
| """Loads the ONNX model once; predict() is safe to call from several threads.""" | |
| def __init__(self, model_path: str | Path, threads: int | None = None): | |
| self.model_path = Path(model_path) | |
| opts = ort.SessionOptions() | |
| opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL | |
| if threads: | |
| opts.intra_op_num_threads = threads | |
| opts.inter_op_num_threads = 1 | |
| self.session = ort.InferenceSession(str(self.model_path), opts, providers=["CPUExecutionProvider"]) | |
| inp = self.session.get_inputs()[0] | |
| self.input_name = inp.name | |
| meta = self.session.get_modelmeta().custom_metadata_map | |
| self.imgsz = inp.shape[2] if isinstance(inp.shape[2], int) else ast.literal_eval(meta.get("imgsz", "[1280]"))[0] | |
| names = ast.literal_eval(meta["names"]) if "names" in meta else dict(enumerate(FALLBACK_NAMES)) | |
| self.names = {int(k): v for k, v in names.items()} | |
| # --- pipeline steps ------------------------------------------------------- | |
| def preprocess(self, bgr: np.ndarray): | |
| img, r, pad = letterbox(bgr, self.imgsz) | |
| x = np.ascontiguousarray(img[:, :, ::-1].transpose(2, 0, 1), dtype=np.float32)[None] / 255.0 | |
| return x, r, pad | |
| def infer(self, x: np.ndarray) -> np.ndarray: | |
| return self.session.run(None, {self.input_name: x})[0] | |
| def postprocess(self, raw, r, pad, shape, conf=0.25, iou=0.7, max_det=300) -> list[Detection]: | |
| pred = raw[0].T # (anchors, 4 + num_classes) | |
| scores_all = pred[:, 4:] | |
| cls = scores_all.argmax(1) | |
| scores = scores_all[np.arange(len(cls)), cls] | |
| keep = scores > conf | |
| if not keep.any(): | |
| return [] | |
| xywh, cls, scores = pred[keep, :4], cls[keep], scores[keep] | |
| boxes = np.empty_like(xywh) | |
| boxes[:, 0] = xywh[:, 0] - xywh[:, 2] / 2 | |
| boxes[:, 1] = xywh[:, 1] - xywh[:, 3] / 2 | |
| boxes[:, 2] = xywh[:, 0] + xywh[:, 2] / 2 | |
| boxes[:, 3] = xywh[:, 1] + xywh[:, 3] / 2 | |
| idx = nms(boxes + cls[:, None] * MAX_WH, scores, iou)[:max_det] | |
| boxes, cls, scores = boxes[idx], cls[idx], scores[idx] | |
| # Ultralytics scale_boxes: undo padding and resize, then clip to the image. | |
| h, w = shape | |
| pad_x = round((self.imgsz - w * r) / 2 - 0.1) | |
| pad_y = round((self.imgsz - h * r) / 2 - 0.1) | |
| boxes = (boxes - [pad_x, pad_y, pad_x, pad_y]) / r | |
| boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, w) | |
| boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, h) | |
| return [ | |
| Detection(int(c), self.names.get(int(c), str(c)), float(s), tuple(float(v) for v in b)) | |
| for b, c, s in zip(boxes, cls, scores) | |
| ] | |
| def predict(self, bgr: np.ndarray, conf: float = 0.25, iou: float = 0.7) -> tuple[list[Detection], Timings]: | |
| t0 = time.perf_counter() | |
| x, r, pad = self.preprocess(bgr) | |
| t1 = time.perf_counter() | |
| raw = self.infer(x) | |
| t2 = time.perf_counter() | |
| dets = self.postprocess(raw, r, pad, bgr.shape[:2], conf, iou) | |
| t3 = time.perf_counter() | |
| return dets, Timings((t1 - t0) * 1000, (t2 - t1) * 1000, (t3 - t2) * 1000) | |
| def draw(bgr: np.ndarray, dets: list[Detection]) -> np.ndarray: | |
| """Return a copy of the image with boxes and labels drawn on it.""" | |
| out = bgr.copy() | |
| thick = max(2, round(max(out.shape[:2]) / 600)) | |
| for d in dets: | |
| x1, y1, x2, y2 = (int(v) for v in d.box) | |
| color = COLORS[d.class_id % len(COLORS)] | |
| cv2.rectangle(out, (x1, y1), (x2, y2), color, thick) | |
| cv2.putText(out, f"{d.class_name} {d.confidence:.2f}", (x1, max(20, y1 - 8)), | |
| cv2.FONT_HERSHEY_SIMPLEX, thick / 2.5, color, thick, cv2.LINE_AA) | |
| return out | |