语义分割只显示指定类别的预测结果

介绍

目前，语义分割的预测结果会将数据集内所有的类别都预测，并展示出来。例如，在下图中我们只想获得人这个类别的预测结果，在不修改模型的前提下，如何实现呢？

方法

可以通过以下代码解决这个问题：

from PIL import Image
import numpy as np
#生成染色版
def get_voc_palette(num_classes):
    n = num_classes
    palette = [0]*(n*3)
    for j in range(0,n):
            lab = j
            palette[j*3+0] = 0
            palette[j*3+1] = 0
            palette[j*3+2] = 0
            i = 0
            while (lab > 0):
                    palette[j*3+0] |= (((lab >> 0) & 1) << (7-i))
                    palette[j*3+1] |= (((lab >> 1) & 1) << (7-i))
                    palette[j*3+2] |= (((lab >> 2) & 1) << (7-i))
                    i = i + 1
                    lab >>= 3
    return palette

#染色代码
def colorize_mask(mask, palette):
    zero_pad = 256 * 3 - len(palette)
    for i in range(zero_pad):
                    palette.append(0)
    new_mask = Image.fromarray(mask.astype(np.uint8)).convert('P')
    new_mask.putpalette(palette)
    return new_mask

#生成指定类别的语义分割预测结果
def generate_img_with_choice_class(img_pth:str,classes:list,num_classes:int):
    #传入 图像路径 和 需要预测的类别  总共的类别
    img = Image.open(img_pth)
    img = np.asarray(img)
    f_img = img.copy()
    for idx,c in enumerate(classes):
        f_img[np.where(img==c)] = 0 #将对应位置置零
    f_img = colorize_mask(f_img,get_voc_palette(num_classes)) # 进行染色处理
    f_img.save('process_img.png')

if __name__ == '__main__':
    img_path = r'D:\1Apython\Pycharm_pojie\data_set\VOCdevkit\VOC2012\voc_test_label\2008_005338.png'
    choice_list = [0,1]
    generate_img_with_choice_class(img_path,choice_list,21)

这个过程主要通过generate_img_with_choice_class(img_pth:str,classes:list,num_classes:int)这个函数实现。它的第一个参数是图像的路径。第二个参数是选择的类别，这里要注意一下选择的类别的含义，是我们不需要的类别，也就是需要置为背景的类别。第三参数是数据集的类别数量，这里我们以Pascal voc 2012 数据集为例，它的类别数是21。
当我们要飞机置位背景时，我们传入[0,1]，因为在该数据集中飞机的标签为1。

上述代码仅以标签图像为例，如果需要获取模型的预测结果，那么我们将读入图片的代码替换为获取模型预测结果的代码即可。