1. 程式人生 > 其它 >python numpy下實現對影象逐個畫素按顏色賦值

python numpy下實現對影象逐個畫素按顏色賦值

技術標籤:python計算機視覺

文章目錄


在做影象語義分割時常常會遇到標註是畫素顏色值型別,這時候我們需要將影象轉換為單通道影象再針對不同的顏色值賦相應的類別,一個簡單的思路就是對影象進行逐畫素遍歷,但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