Python 如何在不使用 `for` 循环的情况下从分割图中获取一批二进制掩码

假设我有一个维度为(1, 1, 6, 6)的分段图a

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)

如何在不使用for循环的情况下获得每个类的二进制掩码？二进制掩码的维度应该是(4, 1, 6, 6)，目前我正在做这样的事情，我想要它没有for循环的原因是a的维度可能会改变，可能会有更多/更少的类.谢谢.

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