图像分割mask图上画出目标检测框

概述：

最近在做跟语义分割有关的变化检测方向的工作。当遇到模型的输出是一张mask图的情况下，单纯用肉眼看检测的效果可能不是很直观。
这时我们可以利用python脚本，调用opencv的cv2.connectedComponentsWithStats这个库，来获取mask连通图中的左上角和右下角的x，y坐标。最后用cv2.rectangle画出矩形框。
参考：https://zhuanlan.zhihu.com/p/59486758
参考的文章作者没有考虑到mask图存在好几个检测框或者没有检测框的情况，我在这个基础上做了改进，并且支持批量转化。

可视化结果：

代码：

import os
import cv2


def mask_find_bboxs(mask):
    retval, labels, stats, centroids = cv2.connectedComponentsWithStats(mask, connectivity=8)  # connectivity参数的默认值为8
    stats = stats[stats[:, 4].argsort()]
    return stats[:-1]  # 排除最外层的连通图


# 根据mask图来从原图上画出识别框
def mask_bbox():
    mask_path = '/home/xiaopeng/dataset/GD_train/9.5/mask/' #输入的mask图
    save_path = '/home/xiaopeng/dataset/GD_train/9.5/save/' #输出的mask图保存路径
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    files = os.listdir(mask_path)

    for file in files:
        # 获取mask（灰度图）
        mask = cv2.imread(mask_path + file, cv2.COLOR_BGR2GRAY)
        # 转换成二值图
        ret, mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
        bboxs = mask_find_bboxs(mask)
        i = 0  # 保存的是需要画的框的个数
        for b in bboxs:
            x0, y0 = b[0], b[1]
            x1 = b[0] + b[2]
            y1 = b[1] + b[3]
            start_point, end_point = (x0, y0), (x1, y1)
            color = (0, 0, 255)  # Red color in BGR；红色：rgb(255,0,0)
            thickness = 2  # Line thickness
            if (x1 - x0) >= 7:  # 连通区域太小时不画框，忽略不计
                if i == 0:
                    mask_BGR = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)  # 转换为3通道图
                    mask_bboxs = cv2.rectangle(mask_BGR, start_point, end_point, color, thickness)
                    i += 1
                else:
                    mask_bboxs = cv2.rectangle(mask_bboxs, start_point, end_point, color, thickness)
                    i += 1
        if i != 0:  # 假如有框时输出图片
            print('{}框的个数为{}'.format(file, i))
            cv2.imwrite(save_path + file, mask_bboxs)


if __name__ == '__main__':
    mask_bbox()