# How to get a batch of binary masks from segmentation map without using `for` loop

suppose I have a segmentation map `a` with dimension of `(1, 1, 6, 6)`

``````print(a)
array([[[[ 0.,  0.,  0.,  0.,  0.,  0.],
[ 0., 15., 15., 16., 16.,  0.],
[ 0., 15., 15., 16., 16.,  0.],
[ 0., 13., 13.,  9.,  9.,  0.],
[ 0., 13., 13.,  9.,  9.,  0.],
[ 0.,  0.,  0.,  0.,  0.,  0.]]]], dtype=float32)
``````

How can I get the binary masks for each class without using for loop? The binary masks should have a dimension of `(4, 1, 6, 6)`, currently im doing something like this and the reason I want it without `for` loop is that the dimension of a might change and there might be more/less classes. Thanks.

``````a1 = np.where(a == 15, 1, 0)
a2 = np.where(a == 16, 1, 0)
a3 = np.where(a == 13, 1, 0)
a4 = np.where(a == 9, 1, 0)
b = np.concatenate((a1, a2, a3, a4), axis=0)

print(b)
array([[[[0, 0, 0, 0, 0, 0],
[0, 1, 1, 0, 0, 0],
[0, 1, 1, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0]]],

[[[0, 0, 0, 0, 0, 0],
[0, 0, 0, 1, 1, 0],
[0, 0, 0, 1, 1, 0],
[0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0]]],

[[[0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
[0, 1, 1, 0, 0, 0],
[0, 1, 1, 0, 0, 0],
[0, 0, 0, 0, 0, 0]]],

[[[0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
[0, 0, 0, 1, 1, 0],
[0, 0, 0, 1, 1, 0],
[0, 0, 0, 0, 0, 0]]]])
``````

### >Solution :

Use `numpy.unique` and broadcasting:

``````u = np.unique(a)
# array([ 0.,  9., 13., 15., 16.], dtype=float32)

out = (a == u[np.nonzero(u)][:,None,None,None]).astype(int)
``````

Output:

``````array([[[[0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
[0, 0, 0, 1, 1, 0],
[0, 0, 0, 1, 1, 0],
[0, 0, 0, 0, 0, 0]]],

[[[0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
[0, 1, 1, 0, 0, 0],
[0, 1, 1, 0, 0, 0],
[0, 0, 0, 0, 0, 0]]],

[[[0, 0, 0, 0, 0, 0],
[0, 1, 1, 0, 0, 0],
[0, 1, 1, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0]]],

[[[0, 0, 0, 0, 0, 0],
[0, 0, 0, 1, 1, 0],
[0, 0, 0, 1, 1, 0],
[0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0]]]])
``````