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

桃扇骨 2024-03-26 14:01 5阅读 0赞

#### 题目 ####

*  基本思路
 *  代码

## 基本思路 ##

做语义分割的同学们，经常需要在原图上可视化预测结果，我们可以借助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()

![在这里插入图片描述][b95b05beb58f41e0bc92ef129023de8d.png]

[b95b05beb58f41e0bc92ef129023de8d.png]: https://image.dandelioncloud.cn/pgy_files/images/2024/03/26/4b79279b7c62477a96680adfab74b377.png