web_ninja - 1 year ago 60

Python Question

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 Source

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