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?

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]])
``````