语义分割在原图上可视化标签

基本思路

做语义分割的同学们，经常需要在原图上可视化预测结果，我们可以借助matplot这个库来进行可视化。

基本原理是首先在画图上展示原图，然后不要清除画布，接着展示标签，为标签增加一个透明度通道，将标签不为零的部分透明度设为0.5，为零的部分透明度设为0 。

代码

import matplotlib.pyplot as plt

# 定义类别颜色
color_mask = [[0, 255, 0], [255, 255, 0], [255, 0, 0]]
for name in os.listdir(image_path):
    image_name = os.path.join(image_path, name)
    label_name = os.path.join(label_path, name)

    image = cv2.imread(image_name)
    label = cv2.imread(label_name, cv2.IMREAD_GRAYSCALE)
    # label = np.expand_dims(label,2).repeat(3, axis=2)

    image = image[:, :, ::-1]
    plt.axis('off')
    plt.imshow(image)
    
    for i in range(3):
        m = np.array(label)
        # 这里需要改成符合自己的数值，一定要大于分类数，比如有3类，这里可以写4，5...
        m[m != i] = 4
        m[m == i] = 1
        m[m == 4] = 0
        img = np.ones( (m.shape[0], m.shape[1], 3) )
        # 随机生成颜色
        # color = np.random.random((1, 3)).tolist()[0]
        color = color_mask[i]
        for j in range(3):
            img[:,:,j] = color[j]
        # 增加透明度通道，叠加到原图上
        plt.imshow(np.dstack( (img, m*0.5) ))

    # 显示图像
    # plt.show()

    # 保存图像，去除白边
    plt.savefig(os.path.join('E:/baidu/myproject1/figure', name), bbox_inches='tight',pad_inches = -0.1)

    # 清除画布
    plt.clf()