我想用Python将3通道RGB图像转换为索引图像。它用于处理训练深层网络进行语义分割的标签。我的意思是它有一个通道,每个像素都是索引,应该从零开始。当然,它们应该有相同的尺寸。转换基于Python dict中的以下映射:
color2index = {
(255, 255, 255) : 0,
(0, 0, 255) : 1,
(0, 255, 255) : 2,
(0, 255, 0) : 3,
(255, 255, 0) : 4,
(255, 0, 0) : 5
}
我实现了一个天真的函数:
^{pr2}$输入im
是由cv2.imread()
创建的一个numpy数组。但是,这个代码非常慢。
由于im
在numpy数组中,我首先用如下方法尝试了numpy的ufunc
:
RGB2index = np.frompyfunc(lambda x: color2index(tuple(x)))
indices = RGB2index(im)
但是结果证明ufunc
每次只需要一个元素。我不能给函数一次三个参数(RGB值)。在
那么有没有其他的方法来优化呢? 如果存在更有效的数据结构,映射就不必是那样的。我注意到Python dict的访问不会花费太多时间,但是从numpy array转换到tuple(它是散列的)需要花费很多时间。在
附言: 我的一个想法是在CUDA中实现一个内核。但这会更复杂。在
更新1:Dan Mašek's Answer效果不错。但首先我们要把RGB图像转换成灰度。当两种颜色的灰度值相同时,可能会出现问题。在
我把工作代码贴在这里。希望它能帮助别人。在
lut = np.ones(256, dtype=np.uint8) * 255
lut[[255,29,179,150,226,76]] = np.arange(6, dtype=np.uint8)
im_out = cv2.LUT(cv2.cvtColor(im, cv2.COLOR_BGR2GRAY), lut)
这里有一个小的实用函数来转换图像(np.数组)到每像素标签(索引),也可以是一个热编码:
如果提供了
None
,它将计算(并返回)颜色码本。在实际上for循环需要很多时间。在
也许上面的代码可以帮助你
你查过枕头库了吗?我记得,它有一些类和方法来处理颜色转换。请参见:https://pillow.readthedocs.io/en/4.0.x/reference/Image.html#PIL.Image.Image.convert
相关问题 更多 >
编程相关推荐