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

Questions

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

May 26, 2023

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