python 将numpy,tensor 转成“索引图“

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]
.....
]