OpenCV4.5.4+YOLOv57.0分割推理演示

原创 OpenCV学堂 2022-11-30 17:29

点击上方蓝字关注我们

微信公众号:OpenCV学堂

关注获取更多计算机视觉与深度学习知识

昨天发了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返回一个,这边是返回两个,所以需要修改一下代码把代码从:
def detect(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()
    return preds

修改为:

def detect(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((201))
    input_blob = np.expand_dims(input_x, 0)
    net.setInput(input_blob)
    layer = net.getUnconnectedOutLayersNames()
    masks, preds = net.forward(layer)
    return preds, 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)
for i in range(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)), (160160))

    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(0256), np.random.randint(0256), np.random.randint(0256)]
    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, (000))

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

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

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


学习YOLOv5最新版从训练到部署(C++,Python)
扫描下面的视频教程



扫码查看OpenCV+OpenVIO+Pytorch系统化学习路线图

 推荐阅读 

CV全栈开发者说 - 从传统算法到深度学习怎么修炼

2022入坑深度学习,我选择Pytorch框架!

Pytorch轻松实现经典视觉任务

教程推荐 | Pytorch框架CV开发-从入门到实战

OpenCV4 C++学习 必备基础语法知识三

OpenCV4 C++学习 必备基础语法知识二

OpenCV4.5.4 人脸检测+五点landmark新功能测试

OpenCV4.5.4人脸识别详解与代码演示

OpenCV二值图象分析之Blob分析找圆

OpenCV4.5.x DNN + YOLOv5 C++推理

OpenCV4.5.4 直接支持YOLOv5 6.1版本模型推理

OpenVINO2021.4+YOLOX目标检测模型部署测试

比YOLOv5还厉害的YOLOX来了,官方支持OpenVINO推理

OpenCV学堂 专注计算机视觉开发技术分享,技术框架使用,包括OpenCV,Tensorflow,Pytorch教程与案例,相关算法详解,最新CV方向论文,硬核代码干货与代码案例详解!作者在CV工程化方面深度耕耘15年,感谢您的关注!
评论 (0)
我要评论
0
0
点击右上角,分享到朋友圈 我知道啦
请使用浏览器分享功能 我知道啦