这段时间部署Yolov5系列模型，想着先使用ONNXRuntime运行一下转换后的onnx模型，但是看着官方给出的detect.py，我陷入了沉思。

我是真的不想抠代码！于是乎，秉承着程序猿的优良传统，我打开了百度。

结果可想而知，不然你也不会看到我的文章了。草草看了一下，大部分是就给出一部分代码，然后告诉你后处理的部分源码里都有............

拜托，我就是不想抠源代码才来百度的啊，哎，还是要靠自己抠代码啊。

2023.04.23--完善了代码后处理部分，使其能够应用到任何一张图片。

yolov5官方代码，链接如下 https://github.com/ultralytics/yolov5.git

一. pytorch模型转onnx模型

什么，你不会导出onnx模型？那就看看官方集成的export.py吧

Yolo官方给出的exporty.py确实很好用，在命令行运行如下代码就能轻松的导出onnx模型。

python3 export.py --weights=./demo.pt --include=onnx --img=640 --batch=1 --opset=12 --simplify

当然，你也可以参看我的另一篇文章，学习并尝试如何转化onnx模型：ONNX系列一：ONNX的使用,从转化到推理

参数介绍如下：

--weights	原pytorch模型
--include	转换的目标格式，支持onnx, torchscript等
--img	模型的输入大小，默认是640
--batch	batch size
--opset	onnx算子集版本
--simplify	简化onnx模型
--half	生成fp16模型，使用该参数时需要指定cuda device，eg：--device=0

二.使用ONNXRuntime进行推理

1.环境要求

python3.7

2.代码实现

okok，能抠的不能抠的都给你抠出来了，你只要动动小手改一下main函数中的路径就能用了!

import cv2
import numpy as np
import onnxruntime as rtCLASSES = {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, eps=1e-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), auto=True, scaleFill=False, scaleup=True, stride=32):# Resize and pad image while meeting stride-multiple constraintsshape = 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 paddingratio = r, r  # width, height ratiosnew_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 paddingif auto:  # minimum rectangledw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh paddingelif scaleFill:  # stretchdw, dh = 0.0, 0.0new_unpad = (new_shape[1], new_shape[0])ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratiosdw /= 2  # divide padding into 2 sidesdh /= 2if shape[::-1] != new_unpad:  # resizeim = cv2.resize(im, new_unpad, 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))im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add borderreturn im, ratio, (dw, dh)def onnx_inf(onnxModulePath, data):sess = rt.InferenceSession(onnxModulePath)input_name = sess.get_inputs()[0].nameoutput_name = sess.get_outputs()[0].namepred_onnx = sess.run([output_name], {input_name: data.reshape(1, 3, 640, 640).astype(np.float32)})return pred_onnxdef xywh2xyxy(x):# Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right# isinstance 用来判断某个变量是否属于某种类型y = np.copy(x)y[..., 0] = x[..., 0] - x[..., 2] / 2  # top left xy[..., 1] = x[..., 1] - x[..., 3] / 2  # top left yy[..., 2] = x[..., 0] + x[..., 2] / 2  # bottom right xy[..., 3] = x[..., 1] + x[..., 3] / 2  # bottom right yreturn ydef nms_boxes(boxes, scores):x = boxes[:, 0]y = boxes[:, 1]w = boxes[:, 2] - boxes[:, 0]h = boxes[:, 3] - boxes[:, 1]areas = w * horder = 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 * h1ovr = inter / (areas[i] + areas[order[1:]] - inter)inds = np.where(ovr <= 0.45)[0]order = order[inds + 1]keep = np.array(keep)return keepdef non_max_suppression(prediction,conf_thres=0.25,iou_thres=0.45,classes=None,agnostic=False,multi_label=False,labels=(),max_det=300,nm=0,  # number of masks
):"""Non-Maximum Suppression (NMS) on inference results to reject overlapping detectionsReturns:list of detections, on (n,6) tensor per image [xyxy, conf, cls]"""# Checksassert 0 <= conf_thres <= 1, f'Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0'assert 0 <= iou_thres <= 1, f'Invalid 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 outputbs = prediction.shape[0]  # batch sizenc = prediction.shape[2] - nm - 5  # number of classesxc = prediction[..., 4] > conf_thres  # candidates# Settingsmax_wh = 7680  # (pixels) maximum box width and heightmax_nms = 30000  # maximum number of boxes into torchvision.ops.nms()redundant = True  # require redundant detectionsmulti_label &= nc > 1  # multiple labels per box (adds 0.5ms/img)merge = False  # use merge-NMSmi = 5 + nc  # mask start indexoutput = [np.zeros((0, 6 + nm))] * bsfor xi, x in enumerate(prediction):  # image index, image inferencex = x[xc[xi]]  # confidenceif labels and len(labels[xi]):lb = labels[xi]v = np.zeros(len(lb), nc + nm + 5)v[:, :4] = lb[:, 1:5]  # boxv[:, 4] = 1.0  # confv[range(len(lb)), lb[:, 0].long() + 5] = 1.0  # clsx = np.concatenate((x, v), 0)# If none remain process next imageif not x.shape[0]:continuex[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf# Box/Maskbox = 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_tuple=False).Tx = np.concatenate((box[i], x[i, 5 + j, None], j[:, None].float(), mask[i]), 1)else:  # best class onlyconf = 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 classif classes is not None:x = x[(x[:, 5:6] == np.array(classes, device=x.device)).any(1)]# Check shapen = x.shape[0]  # number of boxesif not n:  # no boxescontinueindex = x[:, 4].argsort(axis=0)[:max_nms][::-1]x = x[index]# Batched NMSc = x[:, 5:6] * (0 if agnostic else max_wh)  # classesboxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scoresi = 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 matrixweights = iou * scores[None]  # box weightsx[i, :4] = np.multiply(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # merged boxesif redundant:i = i[iou.sum(1) > 1]  # require redundancyoutput[xi] = x[i]return outputdef clip_boxes(boxes, shape):# Clip boxes (xyxy) to image shape (height, width)boxes[..., [0, 2]] = boxes[..., [0, 2]].clip(0, shape[1])  # x1, x2boxes[..., [1, 3]] = boxes[..., [1, 3]].clip(0, shape[0])  # y1, y2def scale_boxes(img1_shape, boxes, img0_shape, ratio_pad=None):# Rescale boxes (xyxy) from img1_shape to img0_shapeif ratio_pad is None:  # calculate from img0_shapegain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  # gain  = old / newpad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2  # wh paddingelse:gain = ratio_pad[0][0]pad = ratio_pad[1]boxes[..., [0, 2]] -= pad[0]  # x paddingboxes[..., [1, 3]] -= pad[1]  # y paddingboxes[..., :4] /= gainclip_boxes(boxes, img0_shape)return boxesif __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))# preprocessim = letterbox(img, imgsz, auto=True)[0]  # padded resizeim = im.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGBim = np.ascontiguousarray(im)  # contiguousim = im.astype(np.float32)im /= 255  # 0 - 255 to 0.0 - 1.0if len(im.shape) == 3:im = im[None]  # expand for batch dim# inferencepred = onnx_inf(onnxModulePath, im)# NMSconf_thres = 0.25  # confidence thresholdiou_thres = 0.45  # NMS IOU thresholdmax_det = 1000  # maximum detections per imageclasses = None  # filter by class: --class 0, or --class 0 2 3agnostic_nms = False  # class-agnostic NMSpred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)# Process predictionsseen = 0for i, det in enumerate(pred):  # per imageseen += 1if len(det):# Rescale boxes from img_size to im0 sizedet[:, :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, decimals=2)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)

运行结果如下：

Yolov5在ONNXRuntime上的推理实现相关推荐

【CVPR2020-中科院计算所】多模态GNN：在视觉信息和场景文字上联合推理
关注上方"深度学习技术前沿",选择"星标公众号", 资源干货,第一时间送达! 即使有可靠的OCR模型,要回答需要在图片中阅读文字的问题,也对现有模型构成了一个挑 ...
【转】知识图谱上推荐推理的模仿学习框架
///又是一篇完全看不懂的文章,泪目了.唯叹一声,道路阻且长--- 原文标题:"SIGIR 2020 | 知识图谱上推荐推理的模仿学习框架" 原文地址:https://www.ms ...
基于PP-ShiTu的零售商品结算系统设计与在Jetson Nano上部署推理使用
文章目录第一章作品概述 1.1 团队介绍 1.2 背景 1.3 简介 1.4 创新点与问题 1.5 发展前景 1.6 技术路线第2章技术方案 2.1 主体检测 2.2 特征提取 2.3 向量检 ...
在英特尔® 硬件上加快推理速度的几个步骤
为了支持云开发人员从云端到边缘的旅程,我们构建了多个加速器.我们将在本博文中介绍其中三个加速器.您可以使用 AWS SageMaker 在 AWS 云中构建和训练模型,然后使用 OpenVINO™ 工 ...
yolov5 onnx 前后处理+运行推理（暂记）
代码在这个(实例分割)基础上改的,虽然跑通了,还是很混乱,这里先简单记录一下处理的流程: yolov5 环境设置 yolov5 网络结构 ONNX yolov5导出 convert error --g ...
【YOLOv5】LabVIEW+OpenVINO让你的YOLOv5在CPU上飞起来
文章目录前言一.OpenVINO是什么二.LabVIEW视觉工具包下载与配置 1.视觉工具包的下载安装 2.OpenVINO toolkit下载安装三.模型获取四.LabVIEW+OpenV ...
SIGIR 2020 | 知识图谱上推荐推理的模仿学习框架
编者按:尽管知识图谱推理的发展前景广阔,但在收敛性和可解释性上仍存在一定的问题.微软亚洲研究院的研究员利用一个基于元启发式方法的示例路径抽取方法来以较低的标记代价提取示例路径集合,进而提出了一个对抗的 ...
mdp框架_SIGIR 2020 | 知识图谱上推荐推理的模仿学习框架
编者按:尽管知识图谱推理的发展前景广阔,但在收敛性和可解释性上仍存在一定的问题.微软亚洲研究院的研究员利用一个基于元启发式方法的示例路径抽取方法来以较低的标记代价提取示例路径集合,进而提出了一个对抗的 ...
FastFormers 论文解读：可以使Transformer 在CPU上的推理速度提高233倍
自Transformers诞生以来,紧随其后的是BERT,在几乎所有与语言相关的任务中,无论是问题回答,情感分析,文本分类还是文本生成,都占据着NLP的主导地位. 与RNN和LSTM消失的梯度问题(不 ...

Yolov5在ONNXRuntime上的推理实现

目录

一. pytorch模型转onnx模型

二.使用ONNXRuntime进行推理

1.环境要求

2.代码实现

Yolov5在ONNXRuntime上的推理实现相关推荐

最新文章

热门文章