AI模型识别出来的结果要想转成图片, 可以转成RGB格式的3维图, 也可以转成1维的索引图,1维的索引图比较方便后续处理

索引图是 1维的,
RGB是3维的.
网上多数是3维的, 很少有提及1维的. 研究了一天终于搞定了.
下面的代码就是如何将AI模型输出转成1维的索引图.

		 
		...上面是torch识别出来的结果, 这里就不贴了.输出的都在output里面, output 矩阵结构是  BCHW 
        maskimg = torch.max(output[:3], 1)[1].detach().clamp_(0, 255).to('cpu', torch.uint8).numpy() # [1] 是第一张图片的输出结果. 
        maskimg = maskimg.transpose(1,2,0) # 从CHW变成了 HWC  
        maskimg = np.sum(maskimg, axis=2) # sum之后就从 HWC 768*1024*1 变成了 1024*768 格式 这里还可以有其它更合适的方法减维度, 请原谅我不会, 这里临时凑一下了. 反正是1个维度的,  sum后还是本身.  
        # maskimg = np.dstack(maskimg,maskimg,maskimg) 
        print("maskimg.shape:",maskimg.shape) 
        # plt.imshow(maskimg) 
        # plt.show() 
         
        # maskimg = Image.fromarray(maskimg,mode='P')# 转成索引图 事实证明这样是不行的.  
        maskimg = Image.fromarray(maskimg).convert('P') # 转成索引图 
 
		colors=[ 
			[0, 0, 0],    #数值0的色彩 
			[128, 0, 0],  #数值1的色彩 
			[0, 128, 0],  #数值2的色彩 
			[128, 128, 0],#数值3的色彩 
			[0, 0, 128],  #数值4的色彩 
			[128, 0, 128], #数值5的色彩 
			..... #可以随便加, 最多应该是uint8的上限 
		] 
		palette = np.array(colors).reshape(-1).tolist() 
 
        maskimg.putpalette(palette)  # 给每个索引一个色彩, 
         
        maskimg.show() 

索引图占存储空间比较小.
写代码就是 x,y就够了. 不用考虑RGB , 读取出来的数据就是需要使用的具体分类结果数据. 比起RGB要更方便更灵活.

顺便附赠一个生成有规律 不随机色彩的代码.

 
def get_color_map_list(num_classes): 
    """ Returns the color map for visualizing the segmentation mask, 
        which can support arbitrary number of classes. 
    Args: 
        num_classes: Number of classes 
    Returns: 
        The color map 
    """ 
    color_map = num_classes * [0, 0, 0] 
    for i in range(0, num_classes): 
        j = 0 
        lab = i 
        while lab: 
            color_map[i * 3] |= (((lab >> 0) & 1) << (7 - j)) 
            color_map[i * 3 + 1] |= (((lab >> 1) & 1) << (7 - j)) 
            color_map[i * 3 + 2] |= (((lab >> 2) & 1) << (7 - j)) 
            j += 1 
            lab >>= 3 
    color_map = [color_map[i:i + 3] for i in range(0, len(color_map), 3)] 
    return color_map 
 
#  print(get_color_map_list(500)) 
 
# 下面的颜色是由上面的代码生成的500种颜色.# 下面的颜色是由上面的代码生成的500种颜色. 
colors = [ 
[0, 0, 0], 
[128, 0, 0], 
[0, 128, 0], 
[128, 128, 0], 
[0, 0, 128], 
[128, 0, 128], 
[0, 128, 128], 
[128, 128, 128], 
[64, 0, 0], 
[192, 0, 0], 
[64, 128, 0], 
[192, 128, 0], 
[64, 0, 128], 
[192, 0, 128] 
..... 
]