在OpenCV4.5.4版本上完成YOLOv5 7.0推理演示與測(cè)試

昨天發(fā)了YOLOv5 7.0支持實(shí)例分割的推文，收到不少留言問推理速度怎么樣，所以我今天測(cè)試了一下，選擇的是YOLOv5s的SEG模型，導(dǎo)出ONNX格式之后，在OpenCV4.5.4版本上完成了推理演示與測(cè)試。

ONNX格式輸入與輸出

首先需要把yolov5s-seg.pt文件導(dǎo)出為ONNX格式，這個(gè)很簡(jiǎn)單，一條命令行搞定：

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

運(yùn)行結(jié)果如下：

導(dǎo)出之后查看輸入與輸出格式顯示如下：

其中輸入部分跟YOLOv5對(duì)象檢測(cè)沒有什么分別，都是NCHW格式圖像輸入，甚至預(yù)處理都完全一致。

輸出部分內(nèi)容分為兩個(gè)部分，output0主要是box框架信息，跟mask預(yù)測(cè)的1x32個(gè)向量，前面85個(gè)解析跟YOLOv5對(duì)象檢測(cè)完成一致，后面32向量是解析mask的時(shí)候會(huì)使用的。

output1格式是1x32x160x160, 針對(duì)每個(gè)box通過(guò)boxes部分的1x32 跟它點(diǎn)乘機(jī)得到1x160x160 就得到這個(gè)box對(duì)應(yīng)的預(yù)測(cè)mask信息，然后根據(jù)box大小從mask中截取roi之后，疊加到輸出結(jié)果上就可以了。

OpenCV DNN推理

整個(gè)代碼實(shí)現(xiàn)部分絕大部分跟OpenCV DNN部署YOLOv5對(duì)象檢測(cè)一致，需要修改的只有兩個(gè)地方，一個(gè)是推理時(shí)候的預(yù)測(cè)結(jié)果，YOLOv5返回一個(gè)，這邊是返回兩個(gè)，所以需要修改一下代碼把代碼從：

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

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對(duì)象檢測(cè)的NMS之后的解析代碼即可。最終Python版本OpenCV DNN推理的運(yùn)行效果如下：

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

基本上可以跑到40FPS左右，感覺很不錯(cuò)了！

審核編輯：劉清