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