CODEWITHSUNDEEP
pythonnumpyopencvvectorization
Oscar de Leeuw
Oscar de Leeuw

Reputation: 126

Vectorize pixelwise assignment for color mapping of a segmentation mask

I want to color a segmentation mask, so it is viewable by a human.

So I have a mask M[y,x,1], where every element is between 0 and n, where n is the amount of classes in the segmentation mask. Furthermore I have a colour table T[n,1,3] where I map every class to a colour in BGR. Lastly I have my colored mask image O[y,x,3] which should have the colour value (defined in T) of the class (defined in M).

I have solved it pixelwise with the following code:

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

But this implementation is horribly slow. Disregard the hardcoded colour table, this can be extracted later :)

Upvotes: 0

Views: 539

Answers (1)

Oscar de Leeuw
Oscar de Leeuw

Reputation: 126

Solved using cv2.LUT as suggested by @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

Ofcourse you can drag the creation of the colour_table outside of this function and pass it as an argument. Also, make sure that segment_mask is CV_8UC1 or CV_8UC3.

Upvotes: 1

Related Questions