I'm new to python and still trying to find a good way to solve this problem:
I have a 2D array label stores the labels (from 0 to 10; 0 is the background, 1 is for the tree, etc.) of an image. From this array, I want to create an RGB image rgb = np.zeros((height, width , 3), np.uint8) on which each pixel in rgb will have the colour depends on the value of label. For example, rgb[4, 8] = green_color if label[4, 8] = 1. What is the most efficient way to do this?
Assuming you have an array named colors of shape 11 x 3 holding all possible colors for the 11 labels (as label ranges from 0 to 10), here's one approach using integer indexing -
img = colors[label.ravel()].reshape(label.shape+(3,)).astype('uint8')
Sample run to verify output -
In [58]: label
Out[58]:
array([[ 2, 4, 0],
[ 2, 9, 10],
[ 6, 10, 10],
[ 4, 0, 4],
[ 1, 4, 10],
[ 8, 1, 7],
[ 9, 8, 0]])
In [59]: colors
Out[59]:
array([[ 28, 175, 15],
[ 0, 255, 0], # Green color at label = 1
[228, 12, 104],
[198, 135, 99],
[126, 124, 222],
[ 35, 78, 14],
[ 64, 61, 0],
[ 34, 49, 147],
[240, 1, 174],
[252, 1, 181],
[171, 114, 191]])
In [60]: img = colors[label.ravel()].reshape(label.shape+(3,))
In [61]: label[4,0]
Out[61]: 1
In [62]: img[4,0]
Out[62]: array([ 0, 255, 0])
In [63]: label[5,1]
Out[63]: 1
In [64]: img[5,1]
Out[64]: array([ 0, 255, 0])
np.choose. np.take seems to be even faster (see my update).raveling and reshapeing - colors[labels] will work fine (it's more or less equivalent to colors.take(labels, axis=0))
The following works, thanks to boolean indexing:
label = np.asarray(( (0,1,1,1), (2,2,0,1) ))
# label will be:
# array([[0, 1, 1, 1],
# [2, 2, 0, 1]])
# or (0, 255, 0) or so
green_color = np.asarray((0,1,0))
# initialize empty image
img = np.zeros((2,4,3))
# set pixels where label==1 to green
img[label == 1] = green_color
You could also use np.choose:
gen = np.random.RandomState(0)
labels = gen.randint(3, size=(5, 6)) # random labels between 0 and 2
print(repr(labels))
# array([[0, 1, 0, 1, 1, 2],
# [0, 2, 0, 0, 0, 2],
# [1, 2, 2, 0, 1, 1],
# [1, 1, 0, 1, 0, 0],
# [1, 2, 0, 2, 0, 1]])
colors = np.array([[255, 0, 0], # red
[0, 255, 0], # green
[0, 0, 255]]) # blue
rgb = labels[..., None].choose(colors)
# labels[0, 1] == 1, so rgb[0, 1] should be [0, 255, 0] (i.e. green)
print(repr(rgb[0, 1]))
# array([ 0, 255, 0])
Another (much faster!) option would be np.take, e.g.:
rgb = colors.take(labels, axis=0)
This can be done more succinctly (but not quite as quickly) by using labels as an index array:
rgb = colors[labels]
Some benchmarking:
# my original answer
In [1]: %%timeit colors = np.random.randint(256, size=(11, 3)); labels = np.random.randint(11, size=(512, 512))
labels[..., None].choose(colors)
....:
10 loops, best of 3: 52 ms per loop
# Divakar's answer
In [2]: %%timeit colors = np.random.randint(256, size=(11, 3)); labels = np.random.randint(11, size=(512, 512))
colors[labels.ravel()].reshape(labels.shape+(3,))
....:
100 loops, best of 3: 4.44 ms per loop
# using take
In [3]: %%timeit colors = np.random.randint(256, size=(11, 3)); labels = np.random.randint(11, size=(512, 512))
colors.take(labels, axis=0)
....:
The slowest run took 4.96 times longer than the fastest. This could mean that an
intermediate result is being cached
1000 loops, best of 3: 1.25 ms per loop
# using integer array indexing
In [4]: %%timeit colors = np.random.randint(256, size=(11, 3)); labels = np.random.randint(11, size=(512, 512))
....: colors[labels]
....:
100 loops, best of 3: 4.19 ms per loop
take had axis option, really good!take is really fast, faster than indexing itself! Do you know why? Some internal optimizations?take can be simpler due to the fact that it only supports integer index arrays, whereas the normal [...] syntax has to deal with any combination of slice objects, int scalars, boolean index arrays etc. I've never looked into the implementations, though.