onnx快速部署YOLO模型

1、准备和环境

首先需要将yolov5模型训练好的最佳权重文件转化为.onnx格式以备使用。不会的小伙伴可以参考yolov5的官方文档，使用yolov5官方的 export.py 脚本进行转换，或者参考一些博客链接，这里不做详细解析。
基本环境配置，相比于yolov5模型源码部署，使用onnx的方式部署会省下不少配置环境的问题，只需要几个关键的第三方库即可完成。

numpy>=1.22.3
onnxruntime>=1.13.1
Pillow>=9.3.0
python-multipart>=0.0.5
fastapi>=0.88.0
python-multipart>=0.0.5
uvicorn[standard]

2、部署代码

2.1、main.py

主文件主要调用

from PIL import Image,ImageDraw,ImageFont
from utils.operation import YOLO
from utils.Colors import colors
import numpy as np
from matplotlib import pyplot as plt
def draw_anchor(img,det_obj):img = Image.open(img)draw = ImageDraw.Draw(img)font = ImageFont.truetype('arial.ttf', 30)imgw,imgh = img.sizecolors_dt = dict() for i in range(len(det_obj)):if colors_dt.get(det_obj[i]['classes']) is None:colors_dt[det_obj[i]['classes']] = colors(i,True)draw.rectangle(det_obj[i]['crop'],width=3,outline=colors_dt[det_obj[i]['classes']])x1, y1, x2, y2 = tuple(det_obj[i]['crop'])draw.text((x1,y1-35),det_obj[i]['classes'],fill=colors_dt[det_obj[i]['classes']],font=font)imgarr = np.array(img)plt.imshow(imgarr)plt.show()img.show() 
def detect(onnx_path='ReqFile/yolov5s.onnx',img=r'ReqFile/bus.jpg',show=True):'''检测目标，返回目标所在坐标如：{'crop': [57, 390, 207, 882], 'classes': 'person'},...]:param onnx_path:onnx模型路径:param img:检测用的图片:param show:是否展示'''#加载yoloyolo = YOLO(onnx_path=onnx_path)  # 加载yolo类#检测det_obj = yolo.decect(img,conf_thres=0.5, iou_thres=0.25)  # 检测#画锚框draw_anchor(img,det_obj)
if __name__ == '__main__':detect()pass

2.2、operation.py

from io import BytesIOimport onnxruntime
import numpy as np
from PIL import Imagefrom utils.orientation import non_max_suppression, tag_imagesclass ONNXModel(object):def __init__(self, onnx_path):""":param onnx_path:"""self.onnx_session = onnxruntime.InferenceSession(onnx_path)self.input_name = self.get_input_name(self.onnx_session)self.output_name = self.get_output_name(self.onnx_session)def get_output_name(self, onnx_session):"""output_name = onnx_session.get_outputs()[0].name:param onnx_session::return:"""output_name = []for node in onnx_session.get_outputs():output_name.append(node.name)return output_namedef get_input_name(self, onnx_session):"""input_name = onnx_session.get_inputs()[0].name:param onnx_session::return:"""input_name = []for node in onnx_session.get_inputs():input_name.append(node.name)return input_namedef get_input_feed(self, input_name, image_numpy):"""input_feed={self.input_name: image_numpy}:param input_name::param image_numpy::return:"""input_feed = {}for name in input_name:input_feed[name] = image_numpyreturn input_feeddef to_numpy(self, file, shape, gray=False):if isinstance(file, np.ndarray):img = Image.fromarray(file)elif isinstance(file, bytes):img = Image.open(BytesIO(file))passelse:img = Image.open(file)widht, hight = shape# 改变大小 并保证其不失真img = img.convert('RGB')if gray:img = img.convert('L')img = img.resize((widht, hight), Image.ANTIALIAS)# 转换成矩阵image_numpy = np.array(img) # (widht, hight, 3)if gray:image_numpy = np.expand_dims(image_numpy,0)image_numpy = image_numpy.transpose(0, 1, 2)else:image_numpy = image_numpy.transpose(2,0,1) # 转置 (3, widht, hight)image_numpy = np.expand_dims(image_numpy,0)# 数据归一化image_numpy = image_numpy.astype(np.float32) / 255.0return image_numpy
class YOLO(ONNXModel):def __init__(self, onnx_path="ReqFile/yolov5n-7-k5.onnx"):super(YOLO, self).__init__(onnx_path)# 训练所采用的输入图片大小self.img_size = 640self.img_size_h = self.img_size_w = self.img_sizeself.batch_size = 1#数量self.num_classes = 2#标签self.classes = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light','fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow','elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee','skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard','tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple','sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch','potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone','microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear','hair drier', 'toothbrush']def to_numpy(self, file, shape, gray=False):def letterbox_image(image, size):iw, ih = image.sizew, h = sizescale = min(w / iw, h / ih)nw = int(iw * scale)nh = int(ih * scale)image = image.resize((nw, nh), Image.BICUBIC)new_image = Image.new('RGB', size, (128, 128, 128))new_image.paste(image, ((w - nw) // 2, (h - nh) // 2))return new_imageif isinstance(file, np.ndarray):img = Image.fromarray(file)elif isinstance(file, bytes):img = Image.open(BytesIO(file))else:img = Image.open(file)resized = letterbox_image(img, (self.img_size_w, self.img_size_h))img_in = np.transpose(resized, (2, 0, 1)).astype(np.float32)  # HWC -> CHWimg_in = np.expand_dims(img_in, axis=0)img_in /= 255.0return img_indef decect(self, file,conf_thres=0.25, iou_thres=0.45):# 图片转换为矩阵image_numpy = self.to_numpy(file, shape=(self.img_size, self.img_size))input_feed = self.get_input_feed(self.input_name, image_numpy)outputs = self.onnx_session.run(self.output_name, input_feed=input_feed)pred = non_max_suppression(outputs[0],conf_thres, iou_thres)if pred:res = tag_images(np.array(Image.open(file)), pred, self.img_size, self.classes)else:res = []return res

2.3、orientation.py

import time
import numpy as np
#用于控制Python中小数的显示精度
np.set_printoptions(precision=4)
def rescale_boxes(boxes, current_dim, original_shape):""" Rescales bounding boxes to the original shape """orig_h, orig_w = original_shape# The amount of padding that was addedpad_x = max(orig_h - orig_w, 0) * (current_dim / max(original_shape))pad_y = max(orig_w - orig_h, 0) * (current_dim / max(original_shape))# Image height and width after padding is removedunpad_h = current_dim - pad_yunpad_w = current_dim - pad_x# Rescale bounding boxes to dimension of original imageboxes[:, 0] = ((boxes[:, 0] - pad_x // 2) / unpad_w) * orig_wboxes[:, 1] = ((boxes[:, 1] - pad_y // 2) / unpad_h) * orig_hboxes[:, 2] = ((boxes[:, 2] - pad_x // 2) / unpad_w) * orig_wboxes[:, 3] = ((boxes[:, 3] - pad_y // 2) / unpad_h) * orig_hreturn boxes
def tag_images(imgs, img_detections, img_size, classes):imgs = [imgs]results = []for img_i, (img, detections) in enumerate(zip(imgs, img_detections)):# Create plotif detections is not None:# Rescale boxes to original imagedetections = rescale_boxes(detections, img_size, img.shape[:2])for x1, y1, x2, y2, conf, cls_pred in detections:results.append({"crop": [int(i) for i in (x1, y1, x2, y2)],"classes": classes[int(cls_pred)]})else:print("识别失败")return results
# 识别结果解析
def xywh2xyxy(x):# Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] # where xy1=top-left, xy2=bottom-righty = 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 y
def nms(dets, scores, thresh):"""Pure Python NMS baseline."""# x1、y1、x2、y2、以及score赋值x1 = dets[:, 0]  # xminy1 = dets[:, 1]  # yminx2 = dets[:, 2]  # xmaxy2 = dets[:, 3]  # ymaxareas = (x2 - x1 + 1) * (y2 - y1 + 1)# argsort()返回数组值从小到大的索引值order = scores.argsort()[::-1]keep = []while order.size > 0:  # 还有数据i = order[0]keep.append(i)if order.size == 1: break# 计算当前概率最大矩形框与其他矩形框的相交框的坐标xx1 = np.maximum(x1[i], x1[order[1:]])yy1 = np.maximum(y1[i], y1[order[1:]])xx2 = np.minimum(x2[i], x2[order[1:]])yy2 = np.minimum(y2[i], y2[order[1:]])# 计算相交框的面积w = np.maximum(0.0, xx2 - xx1 + 1)h = np.maximum(0.0, yy2 - yy1 + 1)inter = w * h# 计算重叠度IOU：重叠面积/（面积1+面积2-重叠面积）IOU = inter / (areas[i] + areas[order[1:]] - inter)left_index = (np.where(IOU <= thresh))[0]# 将order序列更新，由于前面得到的矩形框索引要比矩形框在原order序列中的索引小1，所以要把这个1加回来order = order[left_index + 1]return np.array(keep)def non_max_suppression(prediction, conf_thres, iou_thres, classes=None, agnostic=False, multi_label=False,labels=()):"""Runs Non-Maximum Suppression (NMS) on inference resultsReturns:list of detections, on (n,6) tensor per image [xyxy, conf, cls]"""nc = prediction.shape[2] - 5xc = prediction[..., 4] > conf_thres  # candidates# Settingsmin_wh, max_wh = 2, 4096  # (pixels) minimum and maximum box width and heightmax_det = 300  # maximum number of detections per imagemax_nms = 30000  # maximum number of boxes into torchvision.ops.nms()time_limit = 10.0  # seconds to quit afterredundant = True  # require redundant detectionsmulti_label &= nc > 1  # multiple labels per box (adds 0.5ms/img)t = time.time()output = [np.zeros((0, 6))] * prediction.shape[0]for xi, x in enumerate(prediction):x = x[xc[xi]]  # confidence# Cat apriori labels if autolabellingif labels and len(labels[xi]):l = labels[xi]v = np.zeros((len(l), nc + 5))v[:, :4] = l[:, 1:5]  # boxv[:, 4] = 1.0  # confv[range(len(l)), l[:, 0].long() + 5] = 1.0  # clsx = np.concatenate((x, v), 0)# If none remain process next imageif not x.shape[0]:continue# Compute confx[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf# Box (center x, center y, width, height) to (x1, y1, x2, y2)box = xywh2xyxy(x[:, :4])if multi_label:i, j = (x[:, 5:] > conf_thres).nonzero()x = np.concatenate((box[i], x[i, j + 5, None], j[:, None]), 1)else:  # best class onlyconf = x[:, 5:].max(1, keepdims=True)j = x[:, 5:].argmax(1)j = np.expand_dims(j, 0).Tx = np.concatenate((box, conf, j), 1)[conf.reshape(1, -1)[0] > conf_thres]# Filter by classif classes is not None:x = x[(x[:, 5:6] == np.array(classes)).any(1)]# Check shapen = x.shape[0]  # number of boxesif not n:  # no boxescontinueelif n > max_nms:  # excess boxesx = x[x[:, 4].argsort(descending=True)[:max_nms]]  # sort by confidence# 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, scores, iou_thres)  # NMSif i.shape[0] > max_det:  # limit detectionsi = i[:max_det]output[xi] = x[i]if (time.time() - t) > time_limit:print(f'WARNING: NMS time limit {time_limit}s exceeded')break  # time limit exceededreturn output