Numpy way to integer-mask an array

Question

I have a multi-class segmentation mask eg.

[1 1 1 2 2 2 2 3 3 3 3 3 3 2 2 2 2 4 4 4 4 4 4]

And going to need to get binary segmentation masks for each value

i.e.

[1 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 1 1 1 1 0 0 0 0 0 0 1 1 1 1 0 0 0 0 0 0]

[0 0 0 0 0 0 0 1 1 1 1 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 1 1 1 1 1 1]

Any elegant numpy way to do this?

Ideally an example, where I can set 0 and 1 to other values, if I have to.

Answer 1

Just do "==" as this

import numpy as np
a = np.array([1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 2, 2, 2, 2, 4, 4, 4, 4, 4, 4])
mask1 = (a==1)*5
mask2 = (a==2)*5
mask3 = (a==3)*5
mask4 = (a==4)*5
for mask in [mask1,mask2,mask3,mask4]:
   print(mask)

This gives

[5 5 5 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
[0 0 0 5 5 5 5 0 0 0 0 0 0 5 5 5 5 0 0 0 0 0 0]
[0 0 0 0 0 0 0 5 5 5 5 5 5 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 5 5 5 5 5 5]

You can manipulate the masks further in the same manner, i. e.

mask1[mask1==0] = 3