python numpy下实现对图像逐个像素按颜色赋值

---

在做图像语义分割时常常会遇到标注是像素颜色值类型,这时候我们需要将图像转换为单通道图像再针对不同的颜色值赋相应的类别,一个简单的思路就是对图像进行逐像素遍历,但python下对多维数组逐个进行for循环是很耗时的一件事.numpy底层对矩阵运算做了很多的优化,这时候擅用numpy本身自带的运算可以极大帮助我们加速这一过程.

方法一 利用Numpy 广播机制

def img_array_to_single_val(arr, color_codes):
    result = numpy.ndarray(shape=arr.shape[:2], dtype=int)
    result[:,:] = -1
    for rgb, idx in color_codes.items():
        result[(arr==rgb).all(2)] = idx
    return result

通过广播对矩阵每个像素是否等于rgb进行判断并进行赋值

方法二 通过将RGB一维数组对比转换为标量对比

当矩阵较大时,为了加速我们还可以将RGB一维数组转换为单个标量,进行标量与标量间的对比,加速对比过程

def img_array_to_single_val(image, color_codes):
    image = image.dot(numpy.array([65536, 256, 1], dtype='int32'))
    result = numpy.ndarray(shape=image.shape, dtype=int)
    result[:,:] = -1
    for rgb, idx in color_codes.items():
        rgb = rgb[0] * 65536 + rgb[1] * 256 + rgb[2]
        result[image ==rgb] = idx
    return result

方法三 通过hash加矩阵运算加速对比(耗内存较大)

如果图片更多,并且矩阵更大时,通过矩阵乘法加hash映射加快对比和赋值过程

color_map = numpy.ndarray(shape=(256*256*256), dtype='int32')
color_map[:] = -1
for rgb, idx in color_codes.items():
    rgb = rgb[0] * 65536 + rgb[1] * 256 + rgb[2]
    color_map[rgb] = idx

def img_array_to_single_val(image, color_map):
    image = image.dot(numpy.array([65536, 256, 1], dtype='int32'))
    return color_map[image]

测试:

在python3,在大小为1024X1024的图像中,做十次针对颜色到类别值的映射,上述三种方法耗时如下

可以看到方法3相比方法2快了不少

参考:

https://stackoverflow.com/questions/33196130/replacing-rgb-values-in-numpy-array-by-integer-is-extremely-slow