这段时间部署Yolov5系列模型想着先使用ONNXRuntime运行一下转换后的onnx模型但是看着官方给出的detect.py我陷入了沉思。我是真的不想抠代码于是乎秉承着程序猿的优良传统我打开了百度。结果可想而知不然你也不会看到我的文章了。草草看了一下大部分是就给出一部分代码然后告诉你后处理的部分源码里都有............拜托我就是不想抠源代码才来百度的啊哎还是要靠自己抠代码啊。2023.04.23--完善了代码后处理部分使其能够应用到任何一张图片。yolov5官方代码链接如下 https://github.com/ultralytics/yolov5.git目录一. pytorch模型转onnx模型二.使用ONNXRuntime进行推理1.环境要求2.代码实现一. pytorch模型转onnx模型什么你不会导出onnx模型那就看看官方集成的export.py吧Yolo官方给出的exporty.py确实很好用在命令行运行如下代码就能轻松的导出onnx模型。python3 export.py --weights./demo.pt --includeonnx --img640 --batch1 --opset12 --simplify当然你也可以参看我的另一篇文章学习并尝试如何转化onnx模型ONNX系列一ONNX的使用,从转化到推理参数介绍如下--weights原pytorch模型--include转换的目标格式支持onnx, torchscript等--img模型的输入大小默认是640--batchbatch size--opsetonnx算子集版本--simplify简化onnx模型--half生成fp16模型使用该参数时需要指定cuda deviceeg--device0二.使用ONNXRuntime进行推理1.环境要求python3.72.代码实现okok能抠的不能抠的都给你抠出来了你只要动动小手改一下main函数中的路径就能用了!import cv2 import numpy as np import onnxruntime as rt CLASSES { 0: person, 1: bicycle, 2: car, 3: motorbike, 4: aeroplane, 5: bus, 6: train, 7: truck, 8: boat, 9: traffic light, 10: fire hydrant, 11: stop sign, 12: parking meter, 13: bench, 14: bird, 15: cat, 16: dog, 17: horse, 18: sheep, 19: cow, 20: elephant, 21: bear, 22: zebra, 23: giraffe, 24: backpack, 25: umbrella, 26: handbag, 27: tie, 28: suitcase, 29: frisbee, 30: skis, 31: snowboard, 32: sports ball, 33: kite, 34: baseball bat, 35: baseball glove, 36: skateboard, 37: surfboard, 38: tennis racket, 39: bottle, 40: wine glass, 41: cup, 42: fork, 43: knife, 44: spoon, 45: bowl, 46: banana, 47: apple, 48: sandwich, 49: orange, 50: broccoli, 51: carrot, 52: hot dog, 53: pizza, 54: donut, 55: cake, 56: chair, 57: sofa, 58: potted plant, 59: bed, 60: dining table, 61: toilet, 62: tvmonitor, 63: laptop, 64: mouse, 65: remote, 66: keyboard, 67: cell phone, 68: microwave, 69: oven, 70: toaster, 71: sink, 72: refrigerator, 73: book, 74: clock, 75: vase, 76: scissors, 77: teddy bear, 78: hair drier, 79: toothbrush } def box_iou(box1, box2, eps1e-7): (a1, a2), (b1, b2) box1.unsqueeze(1).chunk(2, 2), box2.unsqueeze(0).chunk(2, 2) inter (np.min(a2, b2) - np.max(a1, b1)).clamp(0).prod(2) return inter / ((a2 - a1).prod(2) (b2 - b1).prod(2) - inter eps) def letterbox(im, new_shape(640, 640), color(114, 114, 114), autoTrue, scaleFillFalse, scaleupTrue, stride32): # Resize and pad image while meeting stride-multiple constraints shape im.shape[:2] # current shape [height, width] if isinstance(new_shape, int): new_shape (new_shape, new_shape) # Scale ratio (new / old) r min(new_shape[0] / shape[0], new_shape[1] / shape[1]) if not scaleup: # only scale down, do not scale up (for better val mAP) r min(r, 1.0) # Compute padding ratio r, r # width, height ratios new_unpad int(round(shape[1] * r)), int(round(shape[0] * r)) dw, dh new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding if auto: # minimum rectangle dw, dh np.mod(dw, stride), np.mod(dh, stride) # wh padding elif scaleFill: # stretch 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] # width, height ratios dw / 2 # divide padding into 2 sides dh / 2 if shape[::-1] ! new_unpad: # resize im cv2.resize(im, new_unpad, interpolationcv2.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)) im cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, valuecolor) # add border return im, ratio, (dw, dh) def onnx_inf(onnxModulePath, data): sess rt.InferenceSession(onnxModulePath) input_name sess.get_inputs()[0].name output_name sess.get_outputs()[0].name pred_onnx sess.run([output_name], {input_name: data.reshape(1, 3, 640, 640).astype(np.float32)}) return pred_onnx def xywh2xyxy(x): # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1top-left, xy2bottom-right # isinstance 用来判断某个变量是否属于某种类型 y np.copy(x) y[..., 0] x[..., 0] - x[..., 2] / 2 # top left x y[..., 1] x[..., 1] - x[..., 3] / 2 # top left y y[..., 2] x[..., 0] x[..., 2] / 2 # bottom right x y[..., 3] x[..., 1] x[..., 3] / 2 # bottom right y return y def nms_boxes(boxes, scores): x boxes[:, 0] y boxes[:, 1] w boxes[:, 2] - boxes[:, 0] h boxes[:, 3] - boxes[:, 1] areas w * h order scores.argsort()[::-1] keep [] while order.size 0: i order[0] keep.append(i) xx1 np.maximum(x[i], x[order[1:]]) yy1 np.maximum(y[i], y[order[1:]]) xx2 np.minimum(x[i] w[i], x[order[1:]] w[order[1:]]) yy2 np.minimum(y[i] h[i], y[order[1:]] h[order[1:]]) w1 np.maximum(0.0, xx2 - xx1 0.00001) h1 np.maximum(0.0, yy2 - yy1 0.00001) inter w1 * h1 ovr inter / (areas[i] areas[order[1:]] - inter) inds np.where(ovr 0.45)[0] order order[inds 1] keep np.array(keep) return keep def non_max_suppression( prediction, conf_thres0.25, iou_thres0.45, classesNone, agnosticFalse, multi_labelFalse, labels(), max_det300, nm0, # number of masks ): Non-Maximum Suppression (NMS) on inference results to reject overlapping detections Returns: list of detections, on (n,6) tensor per image [xyxy, conf, cls] # Checks assert 0 conf_thres 1, fInvalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0 assert 0 iou_thres 1, fInvalid IoU {iou_thres}, valid values are between 0.0 and 1.0 if isinstance(prediction, (list, tuple)): # YOLOv5 model in validation model, output (inference_out, loss_out) prediction prediction[0] # select only inference output bs prediction.shape[0] # batch size nc prediction.shape[2] - nm - 5 # number of classes xc prediction[..., 4] conf_thres # candidates # Settings max_wh 7680 # (pixels) maximum box width and height max_nms 30000 # maximum number of boxes into torchvision.ops.nms() redundant True # require redundant detections multi_label nc 1 # multiple labels per box (adds 0.5ms/img) merge False # use merge-NMS mi 5 nc # mask start index output [np.zeros((0, 6 nm))] * bs for xi, x in enumerate(prediction): # image index, image inference x x[xc[xi]] # confidence if labels and len(labels[xi]): lb labels[xi] v np.zeros(len(lb), nc nm 5) v[:, :4] lb[:, 1:5] # box v[:, 4] 1.0 # conf v[range(len(lb)), lb[:, 0].long() 5] 1.0 # cls x np.concatenate((x, v), 0) # If none remain process next image if not x.shape[0]: continue x[:, 5:] * x[:, 4:5] # conf obj_conf * cls_conf # Box/Mask box xywh2xyxy(x[:, :4]) # center_x, center_y, width, height) to (x1, y1, x2, y2) mask x[:, mi:] # zero columns if no masks # Detections matrix nx6 (xyxy, conf, cls) if multi_label: i, j (x[:, 5:mi] conf_thres).nonzero(as_tupleFalse).T x np.concatenate((box[i], x[i, 5 j, None], j[:, None].float(), mask[i]), 1) else: # best class only conf np.max(x[:, 5:mi], 1).reshape(box.shape[:1][0], 1) j np.argmax(x[:, 5:mi], 1).reshape(box.shape[:1][0], 1) x np.concatenate((box, conf, j, mask), 1)[conf.reshape(box.shape[:1][0]) conf_thres] # Filter by class if classes is not None: x x[(x[:, 5:6] np.array(classes, devicex.device)).any(1)] # Check shape n x.shape[0] # number of boxes if not n: # no boxes continue index x[:, 4].argsort(axis0)[:max_nms][::-1] x x[index] # Batched NMS c x[:, 5:6] * (0 if agnostic else max_wh) # classes boxes, scores x[:, :4] c, x[:, 4] # boxes (offset by class), scores i nms_boxes(boxes, scores) i i[:max_det] # limit detections # 用来合并框的 if merge and (1 n 3E3): # Merge NMS (boxes merged using weighted mean) iou box_iou(boxes[i], boxes) iou_thres # iou matrix weights iou * scores[None] # box weights x[i, :4] np.multiply(weights, x[:, :4]).float() / weights.sum(1, keepdimTrue) # merged boxes if redundant: i i[iou.sum(1) 1] # require redundancy output[xi] x[i] return output def clip_boxes(boxes, shape): # Clip boxes (xyxy) to image shape (height, width) boxes[..., [0, 2]] boxes[..., [0, 2]].clip(0, shape[1]) # x1, x2 boxes[..., [1, 3]] boxes[..., [1, 3]].clip(0, shape[0]) # y1, y2 def scale_boxes(img1_shape, boxes, img0_shape, ratio_padNone): # Rescale boxes (xyxy) from img1_shape to img0_shape if ratio_pad is None: # calculate from img0_shape gain min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain old / new pad (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding else: gain ratio_pad[0][0] pad ratio_pad[1] boxes[..., [0, 2]] - pad[0] # x padding boxes[..., [1, 3]] - pad[1] # y padding boxes[..., :4] / gain clip_boxes(boxes, img0_shape) return boxes if __name__ __main__: onnxModulePath /PATH_to_Yolov5x.onnx IMG_Path /PATH_to_test.jpg imgsz (640, 640) img cv2.imread(IMG_Path) img cv2.resize(img, (640, 640)) # preprocess im letterbox(img, imgsz, autoTrue)[0] # padded resize im im.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB im np.ascontiguousarray(im) # contiguous im im.astype(np.float32) im / 255 # 0 - 255 to 0.0 - 1.0 if len(im.shape) 3: im im[None] # expand for batch dim # inference pred onnx_inf(onnxModulePath, im) # NMS conf_thres 0.25 # confidence threshold iou_thres 0.45 # NMS IOU threshold max_det 1000 # maximum detections per image classes None # filter by class: --class 0, or --class 0 2 3 agnostic_nms False # class-agnostic NMS pred non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_detmax_det) # Process predictions seen 0 for i, det in enumerate(pred): # per image seen 1 if len(det): # Rescale boxes from img_size to im0 size det[:, :4] scale_boxes(im.shape[2:], det[:, :4], img.shape).round() # print(pred) outputs pred[0][:, :6] if len(outputs[:, 4:] 0): for i in outputs: prob i[4] cls int(i[5]) prob np.around(prob, decimals2) if prob 0.4: all_pred_boxes i[:4] for b in range(len(all_pred_boxes)): x1 int(all_pred_boxes[0]) y1 int(all_pred_boxes[1]) x2 int(all_pred_boxes[2]) y2 int(all_pred_boxes[3]) cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 1) cv2.putText(img, CLASSES[cls] str(prob), (x1, y1), cv2.FONT_HERSHEY_TRIPLEX, 0.8, (0, 255, 0), 1, 4) cv2.imwrite(./data/images/test123456789.png, img)运行结果如下