Векторизованное назначение по пикселям для цветового отображения маски сегментации

Question

Я хочу покрасить маску сегментации, чтобы она была видна человеку.

Итак, у меня есть маска M [y, x, 1], где каждый элемент находится между 0 и n, где nколичество классов в маске сегментацииКроме того, у меня есть таблица цветов T [n, 1,3], где я сопоставляю каждый класс с цветом в BGR. Наконец, у меня есть цветное изображение маски O [y, x, 3], в котором должно быть значение цвета (определено в T) класса (определено в M).

Я решил его пиксельно с помощью следующего кода:

def make_colour_mask(segment_mask):
h = segment_mask.shape[0]
w = segment_mask.shape[1]

colour_mask = cv2.cvtColor(segment_mask, cv2.COLOR_GRAY2BGR)

colour_table = [[0,0,0],[255,255,255],[255,0,0],[0,255,0],[0,0,255],[255,0,255],[255,255,255]]

# loop over the image, pixel by pixel
for y in range(0, h):
    for x in range(0, w):
        colour_mask.itemset((y, x, 0), colour_table[segment_mask.item(y,x)][0])
        colour_mask.itemset((y, x, 1), colour_table[segment_mask.item(y,x)][1])
        colour_mask.itemset((y, x, 2), colour_table[segment_mask.item(y,x)][2])

return colour_mask

Но эта реализация ужасно медленная. Не обращайте внимания на жестко закодированную таблицу цветов, это можно извлечь позже:)

Answer 1

Решено с помощью cv2.LUT, как предложено @ Miki.

def make_colour_mask(segment_mask):
    colour_table = np.zeros((256, 1, 3), dtype=np.uint8)
    colour_table[0] = [0, 0, 0]
    colour_table[1] = [255, 255, 255]
    colour_table[2] = [255, 0, 0]
    colour_table[3] = [0, 255, 0]
    colour_table[4] = [0, 0, 255]
    colour_table[5] = [255 ,0 ,255]
    colour_table[6] = [0, 255 ,255]

    colour_mask = cv2.applyColorMap(segment_mask, colour_table)

    return colour_mask

Конечно, вы можете перетащить создание colour_table за пределы этой функции и передать его в качестве аргумента. Также убедитесь, что segment_mask равно CV_8UC1 или CV_8UC3.