在OpenCV4.5.4版本上完成YOLOv5 7.0推理演示与测试

昨天发了YOLOv5 7.0支持实例分割的推文，收到不少留言问推理速度怎么样，所以我今天测试了一下，选择的是YOLOv5s的SEG模型，导出ONNX格式之后，在OpenCV4.5.4版本上完成了推理演示与测试。

ONNX格式输入与输出

首先需要把yolov5s-seg.pt文件导出为ONNX格式，这个很简单，一条命令行搞定：

python export.py --weights yolov5s-seg.pt --include onnx

运行结果如下：

导出之后查看输入与输出格式显示如下：

其中输入部分跟YOLOv5对象检测没有什么分别，都是NCHW格式图像输入，甚至预处理都完全一致。

输出部分内容分为两个部分，output0主要是box框架信息，跟mask预测的1x32个向量，前面85个解析跟YOLOv5对象检测完成一致，后面32向量是解析mask的时候会使用的。

output1格式是1x32x160x160, 针对每个box通过boxes部分的1x32 跟它点乘机得到1x160x160 就得到这个box对应的预测mask信息，然后根据box大小从mask中截取roi之后，叠加到输出结果上就可以了。

OpenCV DNN推理

整个代码实现部分绝大部分跟OpenCV DNN部署YOLOv5对象检测一致，需要修改的只有两个地方，一个是推理时候的预测结果，YOLOv5返回一个，这边是返回两个，所以需要修改一下代码把代码从：

defdetect(image,net):
#1x3x640x640
blob=cv2.dnn.blobFromImage(image,1/255.0,(INPUT_WIDTH,INPUT_HEIGHT),swapRB=True,crop=False)
net.setInput(blob)
preds=net.forward()
returnpreds

修改为：

defdetect(image,net):
rgb=cv.cvtColor(image,cv.COLOR_BGR2RGB)
input_image=cv.resize(src=rgb,dsize=(INPUT_WIDTH,INPUT_HEIGHT))
blob_img=np.float32(input_image)/255.0
input_x=blob_img.transpose((2,0,1))
input_blob=np.expand_dims(input_x,0)
net.setInput(input_blob)
layer=net.getUnconnectedOutLayersNames()
masks,preds=net.forward(layer)
returnpreds,masks

这样就好啦 第二个改动的地方在后处理部分，如何解析出mask部分，这部分我通过翻看YOLOv5 7.0官方推理演示的源码，它是基于torch实现的，我一通猛改之后改成了基于numpy实现。生成mask的代码如下：

color_mask=np.zeros((fh,fw,3),dtype=np.uint8)
black_mask=np.zeros((fh,fw),dtype=np.float32)
mv=cv.split(color_mask)
foriinrange(len(boxes)):
x1,y1,x2,y2=boxes[i]
x1=max(0,x1)
y1=max(0,y1)
classid=class_ids[i]
m1=masks[i]
mask=np.reshape(sigmoid(np.matmul(m1,mask2)),(160,160))

mx1=max(0,np.int((x1*sx)/x_factor))
mx2=max(0,np.int((x2*sx)/x_factor))
my1=max(0,np.int((y1*sy)/y_factor))
my2=max(0,np.int((y2*sy)/y_factor))
mask_roi=mask[my1:my2,mx1:mx2]

result_mask=cv.resize(mask_roi,(x2-x1,y2-y1))
result_mask[result_mask>0.5]=1.0
result_mask[result_mask<= 0.5] = 0.0
    rh, rw = result_mask.shape
    if (y1+rh) >=fh:
rh=fh-y1
if(x1+rw)>=fw:
rw=fw-x1
black_mask[y1:y1+rh,x1:x1+rw]=result_mask[0:rh,0:rw]
mv[2][black_mask==1],mv[1][black_mask==1],mv[0][black_mask==1]=
[np.random.randint(0,256),np.random.randint(0,256),np.random.randint(0,256)]
color=colors[int(classid)%len(colors)]
cv.rectangle(frame,(x1,y1),(x2,y2),color,2)
cv.rectangle(frame,(x1,y1-20),(x2,y1),color,-1)
cv.putText(frame,class_list[classid],(x1,y1-10),cv.FONT_HERSHEY_SIMPLEX,.5,(0,0,0))

把这段代码放在NMS之后，替换YOLOv5对象检测的NMS之后的解析代码即可。最终Python版本OpenCV DNN推理的运行效果如下：

速度这么慢，怒而改成OpenCV DNN C++推理，N卡加持：

基本上可以跑到40FPS左右，感觉很不错了！

审核编辑：刘清