前言
相必大家在对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为自己设置的阈值,这部分大家可以根据自己实际应用中进行更改
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函数这里大家应该都很熟悉了。
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
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(四个位置信息 + 一个置信度)
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的数组中我们需要遍历计算我们需要的信息,循环前定义空的表用作存储使用
# 记录缩放比
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中的时候也应该按照这个顺序进行填入
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的坐标框
在这里我们对结果进行判断,因为存在输出为空的情况,即:在图像中没有检测出符合要求的目标
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
在这里大家可以对绘制的颜色以及字体相关信息自己按照喜好调整!
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)
入口调用
在这里我们调用定义的函数,输入图像的路径和网络的路径以及标签类别总数
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)
附录:
检测结果信息表:
项目地址:
本文正在参加「金石计划 . 瓜分6万现金大奖」