python - 根据颜色值将多维 Numpy 数组转换为二维数组

Question

我有一张图像，它被读取为形状为 (512,512,3) 的 uint8 数组。 现在我想将此数组转换为形状为 (512,512,1) 的 uint8 数组，其中第三轴中的每个像素值都是从颜色值 [255,0,0 ] 到单个类别标签值 [3]，基于以下颜色/类别编码:

1 : [0, 0, 0], 
 2 : [0, 0, 255], 
 3 : [255, 0, 0], 
 4 : [150, 30, 150], 
 5 : [255, 65, 255], 
 6 : [150, 80, 0], 
 7 : [170, 120, 65], 
 8 : [125, 125, 125], 
 9 : [255, 255, 0], 
 10 : [0, 255, 255], 
 11 : [255, 150, 0], 
 12 : [255, 225, 120], 
 13 : [255, 125, 125], 
 14 : [200, 100, 100], 
 15 : [0, 255, 0], 
 16 : [0, 150, 80], 
 17 : [215, 175, 125], 
 18 : [220, 180, 210], 
 19 : [125, 125, 255]

最有效的方法是什么？我想到了遍历所有类并使用 numpy.where，但这显然很耗时。

Answer 1

您可以使用巨大的查找表。令 cls 为 [[0,0,0], [0,0,255], ...] 的 dtype=np.uint8。

LUT = np.zeros(size=(256,256,256), dtype='u1')
LUT[cls[:,0],cls[:,1],cls[:,2]] = np.arange(cls.shape[1])+1
img_as_cls = LUT[img[...,0],img[...,1], img[...,2]]

此解决方案每个像素的复杂度为 O(1)。它的缓存效率也很高，因为实际使用了 LUT 中的一小部分条目。在我的机器上处理 1000x1000 图像大约需要 10 毫秒。

可以通过将 3 色 channel 转换为 24 位整数来稍微改进解决方案。 这是代码

def scalarize(x):
    # compute x[...,2]*65536+x[...,1]*256+x[...,0] in efficient way
    y = x[...,2].astype('u4')
    y <<= 8
    y +=x[...,1]
    y <<= 8
    y += x[...,0]
    return y
LUT = np.zeros(2**24, dtype='u1')
LUT[scalarize(cls)] = 1 + np.arange(cls.shape[0])
simg = scalarize(img)
img_to_cls = LUT[simg]

优化后处理 1000x1000 图像大约需要 5ms。