前言
相必大家在对yolo项目部署的过程中是否存在过:有没有一个代码能将当前github上的yolov4至v7的onnx推理统一一下。在这里由于官网提供的yolo代码 编写习惯不尽相同,大家的学习成本比较大,在这里我为大家带来的是bubbliiiing的yolo系列代码。 目前经我实际测试验证了yolov4、yolov5、yolov5-6.1、yolov7都可以用这套代码实现onnx部署(在这里不多赘述onnx部署的优点了) 本文以yolov7为例子进行讲解(其他版本的实现思路相同)
部署之前
在进行onnx的部署前我们需要导出onnx,这部分可以参考我github项目中的README.md部分
推理框架
整个opencv推理模块分为四部分:
- Anchors、Stride和置信度等参数配置
- 载入onnx权重文件
- 对图像进行推理得到目标坐标点和置信度、目标ID
- 对推理结果进行在图像上绘制相关信息
参数配置
这一部分实现起来比较简单,但是非常重要!
这里我们对网络目标类别进行载入,className文件格式同官网即可,Anchors和Stride填入相对于的版本的信息即可,后续的Width、Height设置为导出onnx时设置的大小,nmsThreshold、boxThreshold、classThreshold为自己设置的阈值,这部分大家可以根据自己实际应用中进行更改
ini
复制代码
className = list(map(lambda x: x.strip(), open('coco.names', 'r').readlines())) # 可替换自己的类别文件 # 替换对应yolo的Anchors值 netAnchors = np.asarray([[12.0, 16.0, 19.0, 36.0, 40.0, 28.0], [36.0, 75.0, 76.0, 55.0, 72.0, 146.0], [142.0, 110.0, 192.0, 243.0, 459.0, 401.0]], dtype=np.float32) netStride = np.asarray([8.0, 16.0, 32.0], dtype=np.float32) netWidth = 640 netHeight = 640 nmsThreshold = 0.80 boxThreshold = 0.80 classThreshold = 0.80
在这里初始化相关配置的时候,也需要生成我们绘制目标框时框的颜色,这里的color_num为自己类别总数。Sigmoid函数这里大家应该都很熟悉了。
ini
复制代码
def GetColors(color_num): ColorList = [] for num in range(color_num): R = random.randint(100, 255) G = random.randint(100, 255) B = random.randint(100, 255) BGR = (B, G, R) ColorList.append(BGR) return ColorList def Sigmoid(x): x = float(x) out = (float(1.) / (float(1.) + exp(-x))) return float(out)
读取ONNX
大家在这里的读取onnx时候我踩到的一个最大的坑就是直接pip install opencv-python。直接pip得到的opencv的版本为4.6.0,当了两天的小白鼠实践最新版本的 opencv,总之计算的结果是莫名其妙。这里推荐大家使用pip install opencv-python==4.5.5.62
在这里定义了一个调用onnx函数,输入onnx网络的路径以及是否调用GPU模块,函数返回值为net
python
复制代码
def readModel(netPath, isCuda=False): try: net = cv2.dnn.readNetFromONNX(netPath) except: return False if isCuda: # GPU net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA) net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA) else: # CPU net.setPreferableBackend(cv2.dnn.DNN_BACKEND_DEFAULT) net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU) return net
推理模块
由于整个推理模块比较多,在这里我为大家分成多个部分进行讲解
图像输入及输出
这里定义推理函数,输入分别是输入图像、读取onnx的网络,网络设点给的宽以及高。由于需要对图像进行推理,在使用opencv读取图像的通道格式为BGR,以及图像维度需要转换为为(1,3,640,640),这里我们需要借助 cv2.dnn.blobFromImage 函数进行快速转换(当大家进行批量读取的时候可以更换为cv2.dnn.blobFromImages)
对转换好的图像进行网络推理并得到输出,这里我们得到的输出为:(1,25200,85)
ps: 25200 = 3 * 80 * 80 * 85 + 3 * 40 * 40 * 85 + 3 * 20 * 20 * 85
80 = 640 / 8
40 = 640 / 16
20 = 640 / 32
85 = 80(类别总数) + 5(四个位置信息 + 一个置信度)
ini
复制代码
def Detect(SrcImg, net, netWidth, netHeight): netInputImg = SrcImg blob = cv2.dnn.blobFromImage(netInputImg, scalefactor=1 / 255.0, size=(netWidth, netHeight), mean=[104, 117, 123], swapRB=True, crop=False) net.setInput(blob) # 把图像放入网络 netOutputImg = net.forward(net.getUnconnectedOutLayersNames()) # 获取输出结果集 netOutputImg = np.array(netOutputImg, dtype=np.float32) pdata = netOutputImg[0]
计算网络输出结果
上述我们对图像进行了归一化操作,那么我们需要计算图像的缩放比;在一个25200 * 85的数组中我们需要遍历计算我们需要的信息,循环前定义空的表用作存储使用
ini
复制代码
# 记录缩放比 ratio_h = float(netInputImg.shape[1] / netHeight) ratio_w = float(netInputImg.shape[0] / netWidth) classIds = [] # 定义结果表 confidences = [] # 定义置信度表 boxes = [] # 定义坐标表 count = 0 # 标记符
我们根据源码的输出计算可以有如下for循环:
- 首先是通过三个Stride值计算得到grid值(80 40 20 的由来)
- 然后是对Anchors的值进行提取为了坐标的计算
- 对每一次的grid进行循环计算
对640 * 640的图像我们总共遍历的次数是25200次,也就对应了25200行的输出信息。我们需要对每一行的box阈值进行计算,对符合要求的box阈值 再进行提取坐标以及类别最大可能值信息进计算。
大家需要注意下,由于我们提取了很多个坐标框,如果这样我们就进行绘制在图像上的话会十分糟糕的,不够优雅,这里我们就需要调用cv2.dnn.NMSBoxes 模块,由于cv2.dnn.NMSBoxes模块里对坐标进行抑制的话坐标的排列顺序为:BOX = (left, top, W, H) ,那么我们在写入BOX中的时候也应该按照这个顺序进行填入
ini
复制代码
for stride in range(3): # netStride = {8.0, 16.0, 32.0} = 3 grid_x = netWidth / netStride[stride] grid_y = netHeight / netStride[stride] grid_x, grid_y = int(grid_x), int(grid_y) # 系统默认是float32,这里是为了下面的循环转为int for anchor in range(3): # netAnchors 的层数 = 3 anchor_w = netAnchors[stride][anchor * 2] anchor_h = netAnchors[stride][anchor * 2 + 1] anchor_w, anchor_h = float(anchor_w), float(anchor_h) for i in range(grid_x): for j in range(grid_y): # 到这的下一行总运行次数是25200 = (80*80*3) + (40*40*3) + (20*20*3) pdatabox = pdata[0][count][4] box_score = Sigmoid(pdatabox) # 获取每一行的box框中含有某个物体的概率 if box_score > boxThreshold: # box的阈值起作用了 scores = pdata[0][count][5:] # 这里的scores理应是一个多维矩阵 _, max_class_socre, _, classIdPoint = cv2.minMaxLoc(scores) # 求最大值以及最大值的位置&位置是元组 max_class_socre = np.asarray(max_class_socre, dtype=np.float64) max_class_socre = Sigmoid(max_class_socre) if max_class_socre > classThreshold: # 类别的置信度起作用 # rect[x,y,w,h] pdatax=pdata[0][count][0] x = (Sigmoid(pdatax) * float(2.) - float(0.5) + j) * netStride[stride] # x pdatay = np.asarray(pdata[0][count][1], dtype=np.float64) y = (Sigmoid(pdatay) * float(2.) - float(0.5) + i) * netStride[stride] # y pdataw=pdata[0][count][2] w = pow(Sigmoid(pdataw) * float(2.), float(2.0)) * anchor_w # w pdatah = pdata[0][count][3] h = pow(Sigmoid(pdatah) * float(2.), float(2.0)) * anchor_h # h left = (x - 0.5 * w) * ratio_w top = (y - 0.5 * h) * ratio_h left, top, W, H = int(left), int(top), int(w * ratio_w), int(h * ratio_h) # 对classIds & confidences & boxes classIds.append(classIdPoint[1]) # 获取最大值的位置 confidences.append(max_class_socre * box_score) boxes.append((left, top, W, H)) count += 1
极大值抑制
这里我们需要对25200次循环的结果进行极大值抑制,参考opencv官网我们需要输入boxes, confidences, classThreshold, nmsThreshold 最后得到的结果为:
- 返回结果的ID(便于从className中找到对应的类别名称);
- 该ID的置信度
- 该ID的坐标框 在这里我们对结果进行判断,因为存在输出为空的情况,即:在图像中没有检测出符合要求的目标
ini
复制代码
nms_result = cv2.dnn.NMSBoxes(boxes, confidences, classThreshold, nmsThreshold) # 抑制处理 返回的是一个数组 result_id, result_confidence, result_boxs = [], [], [] for idx in nms_result: result_id.append(classIds[idx]) result_confidence.append(confidences[idx]) result_boxs.append(boxes[idx]) if len(result_id) == 0: return False else: return result_id, result_confidence, result_boxs
结果绘制
我们终终于将结果获取到了我们需要的结果,这里就需要进行绘制了。在对BOX的结果进行分解: BOX = (left, top, W, H),在绘制的时候需要用到cv2.rectangle和cv2.putText,那么我们需要一张图像的Xmin Xmax Ymin Ymax 其中:Xmin = left Ymin = top Xmax = left + W Ymax = top + H 在这里大家可以对绘制的颜色以及字体相关信息自己按照喜好调整!
ini
复制代码
def drawPred(img, result_id, result_confidence, result_boxs, color): for i in range(len(result_id)): class_id = result_id[i] confidence = round(result_confidence[i], 2) # 保留两位有效数字 box = result_boxs[i] pt1 = (box[0], box[1]) pt2 = (box[0] + box[2], box[1] + box[3]) # 绘制图像目标位置 cv2.rectangle(img, pt1, pt2, color[i], 2, 2) # x1, y1, x2, y2 = box[0], box[1], box[0] + box[2], box[1] + box[3] cv2.rectangle(img, (box[0], box[1]-18), (box[0] + box[2], box[1]), color[i], -1) label = "%s:%s" % (className[class_id], confidence) # 给目标进行添加类别名称以及置信度 FONT_FACE = cv2.FONT_HERSHEY_SIMPLEX cv2.putText(img, label, (box[0] - 2, box[1] - 5), FONT_FACE, 0.5, (0, 0, 0), 1) cv2.imwrite("my640.jpg", img) cv2.imshow("outPut", img) cv2.waitKey(0)
入口调用
在这里我们调用定义的函数,输入图像的路径和网络的路径以及标签类别总数
ini
复制代码
if __name__ == "__main__": img_path = "bus.jpg" model_path = "yolov7.onnx" classNum = 80 color = GetColors(classNum) Mynet = readModel(model_path, isCuda=False) img = cv2.imread(img_path) result_id, result_confidence, result_boxs = Detect(img, Mynet, 640, 640) drawPred(img, result_id, result_confidence, result_boxs, color)
附录:
检测结果信息表:
项目地址: