from typing import List, Dict import cv2 import numpy as np import onnxruntime as ort def letterbox( img: np.ndarray, new_shape=(960, 960), color=(114, 114, 114), auto=False, scale_fill=False, scale_up=True, stride=32, ): shape = img.shape[:2] # h, w if isinstance(new_shape, int): new_shape = (new_shape, new_shape) r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) if not scale_up: r = min(r, 1.0) ratio = (r, r) new_unpad = (int(round(shape[1] * r)), int(round(shape[0] * r))) dw = new_shape[1] - new_unpad[0] dh = new_shape[0] - new_unpad[1] if auto: dw, dh = np.mod(dw, stride), np.mod(dh, stride) elif scale_fill: dw, dh = 0.0, 0.0 new_unpad = (new_shape[1], new_shape[0]) ratio = (new_shape[1] / shape[1], new_shape[0] / shape[0]) dw /= 2 dh /= 2 if shape[::-1] != new_unpad: img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) top = int(round(dh - 0.1)) bottom = int(round(dh + 0.1)) left = int(round(dw - 0.1)) right = int(round(dw + 0.1)) img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) return img, ratio, (dw, dh) def xywh2xyxy(x: np.ndarray) -> np.ndarray: y = np.copy(x) y[:, 0] = x[:, 0] - x[:, 2] / 2 y[:, 1] = x[:, 1] - x[:, 3] / 2 y[:, 2] = x[:, 0] + x[:, 2] / 2 y[:, 3] = x[:, 1] + x[:, 3] / 2 return y def box_iou(box: np.ndarray, boxes: np.ndarray) -> np.ndarray: x1 = np.maximum(box[0], boxes[:, 0]) y1 = np.maximum(box[1], boxes[:, 1]) x2 = np.minimum(box[2], boxes[:, 2]) y2 = np.minimum(box[3], boxes[:, 3]) inter = np.maximum(0, x2 - x1) * np.maximum(0, y2 - y1) area1 = max(0, box[2] - box[0]) * max(0, box[3] - box[1]) area2 = np.maximum(0, boxes[:, 2] - boxes[:, 0]) * np.maximum(0, boxes[:, 3] - boxes[:, 1]) union = area1 + area2 - inter + 1e-7 return inter / union def nms_numpy(boxes: np.ndarray, scores: np.ndarray, iou_thres: float) -> List[int]: if len(boxes) == 0: return [] order = scores.argsort()[::-1] keep = [] while order.size > 0: i = order[0] keep.append(int(i)) if order.size == 1: break ious = box_iou(boxes[i], boxes[order[1:]]) inds = np.where(ious <= iou_thres)[0] order = order[inds + 1] return keep def scale_boxes_from_letterbox(boxes: np.ndarray, ratio, pad, original_shape) -> np.ndarray: # boxes: xyxy in letterboxed image gain = ratio[0] pad_x, pad_y = pad boxes[:, [0, 2]] -= pad_x boxes[:, [1, 3]] -= pad_y boxes[:, :4] /= gain h, w = original_shape[:2] boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, w - 1) boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, h - 1) return boxes class YOLOv5ONNXDetector: def __init__( self, model_path: str, class_names: List[str], img_size: int = 960, device: str = "cpu", ): self.model_path = model_path self.class_names = class_names self.img_size = img_size providers = ["CPUExecutionProvider"] if device.lower() in ["cuda", "gpu"]: providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] self.session = ort.InferenceSession(model_path, providers=providers) active_providers = self.session.get_providers() print("[YOLO] providers:", active_providers) if ( device.lower() in ["cuda", "gpu"] and "CUDAExecutionProvider" not in active_providers ): raise RuntimeError( "请求使用 CUDA 推理，但 YOLO 会话未启用 " "CUDAExecutionProvider。请安装匹配 CUDA/cuDNN 的 " "onnxruntime-gpu。" ) self.input_name = self.session.get_inputs()[0].name self.output_names = [o.name for o in self.session.get_outputs()] def preprocess(self, img_bgr: np.ndarray): img_lb, ratio, pad = letterbox(img_bgr, new_shape=(self.img_size, self.img_size), auto=False) img_rgb = cv2.cvtColor(img_lb, cv2.COLOR_BGR2RGB) img = img_rgb.transpose(2, 0, 1).astype(np.float32) / 255.0 img = np.expand_dims(img, axis=0) return img, ratio, pad def infer( self, img_bgr: np.ndarray, conf_thres: float = 0.25, iou_thres: float = 0.45, max_det: int = 50, ) -> List[Dict]: input_tensor, ratio, pad = self.preprocess(img_bgr) outputs = self.session.run(self.output_names, {self.input_name: input_tensor}) pred = outputs[0] # 常见 YOLOv5 ONNX 输出: [1, num_boxes, 5 + num_classes] if isinstance(pred, list): pred = pred[0] pred = np.asarray(pred) if pred.ndim == 3: pred = pred[0] elif pred.ndim == 2: pass else: raise ValueError(f"不支持的 YOLO 输出维度: {pred.shape}") if pred.shape[1] < 6: raise ValueError(f"YOLO 输出最后一维应至少为 6，当前为: {pred.shape}") boxes_xywh = pred[:, :4] obj_conf = pred[:, 4] if pred.shape[1] == 6: # 某些导出格式可能已经是 x,y,w,h,conf,class_id cls_ids = pred[:, 5].astype(np.int32) scores = obj_conf else: cls_scores = pred[:, 5:] cls_ids = np.argmax(cls_scores, axis=1) cls_conf = cls_scores[np.arange(len(cls_scores)), cls_ids] scores = obj_conf * cls_conf mask = scores >= conf_thres if not np.any(mask): return [] boxes_xywh = boxes_xywh[mask] scores = scores[mask] cls_ids = cls_ids[mask] boxes = xywh2xyxy(boxes_xywh) boxes = scale_boxes_from_letterbox(boxes, ratio, pad, img_bgr.shape) final_results = [] for class_id in np.unique(cls_ids): idxs = np.where(cls_ids == class_id)[0] keep = nms_numpy(boxes[idxs], scores[idxs], iou_thres) for k in keep: original_idx = idxs[k] cid = int(class_id) cname = self.class_names[cid] if 0 <= cid < len(self.class_names) else str(cid) final_results.append({ "bbox_xyxy": boxes[original_idx].astype(float).tolist(), "det_confidence": float(scores[original_idx]), "det_class_id": cid, "det_class_name": cname, }) final_results.sort(key=lambda x: x["det_confidence"], reverse=True) return final_results[:max_det]