web_ninja web_ninja - 4 months ago 9
Python Question

How to speed up the creation of a label tensor from label map in Numpy?

Given a label map of dimensions W X H where each element can take values from {0,..,K-1} I want to output a label tensor of dimensions K X W x H where each element in the K'th map is 1 only if the corresponding value in the labelmap was K. Currently my implementation uses two for loops and is very slow.

p_label = Labelmap with one channel

label = np.zeros((K,p_label.shape[0], p_label.shape[1]))
for i in xrange(p_label.shape[0]):
for j in xrange(p_label.shape[1]):
label[p_label[i,j],i,j] = 1


Is there a better way to do this operation in Numpy using broadcasting?

Answer

You can use the == operator with broadcasting.

For example,

In [19]: W = 5

In [20]: H = 8

In [21]: K = 10

Create a p_label for the example:

In [22]: p_label = np.random.randint(0, K, size=(W, H))

kvals is simply an array containing [0, 1, ..., K-1]:

In [23]: kvals = np.arange(K)

kvals.reshape(-1, 1, 1) converts kvals to an array with shape (K, 1, 1). This is compared using == to p_label. Broadcasting applies, so the result of the comparison has shape (K, W, H). It is a boolean array of the values that you want. .astype(int) converts the result to an integer array. (You can remove that if a boolean array would work for you.)

In [24]: label = (p_label == kvals.reshape(-1, 1, 1)).astype(int)

Here's the original p_label. Note, for example, the locations of the value 0:

In [25]: p_label
Out[25]: 
array([[3, 3, 2, 6, 2, 2, 9, 3],
       [1, 8, 1, 1, 4, 3, 7, 8],
       [5, 9, 1, 0, 7, 2, 8, 0],
       [1, 3, 5, 4, 6, 0, 9, 5],
       [5, 7, 2, 0, 6, 4, 5, 3]])

label[0] is 1 in the positions where p_label is 0.

In [26]: label[0]
Out[26]: 
array([[0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 1, 0, 0, 0, 1],
       [0, 0, 0, 0, 0, 1, 0, 0],
       [0, 0, 0, 1, 0, 0, 0, 0]])
Comments