时间戳

首页 > 解决方案 > 根据颜色值将多维 Numpy 数组转换为二维数组

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

问题描述

我有一个图像，它被读取为形状为 的 uint8 数组(512,512,3)。现在我想将此数组转换为 shape 的 uint8 数组(512,512,1)，其中第三轴中的每个像素值都根据以下颜色/类编码从颜色值转换[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，但这显然很耗时。

标签： pythonarraysimagenumpyindexing

解决方案

你可以使用巨大的查找表。让 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 色通道转换为 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。

推荐阅读