OpenCV4.5.4+YOLOv5 7.0分割推理演示

微信公众号：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((2, 0, 1))
    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)), (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推理的运行效果如下：