如何不使用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)。目前我采用逐个判断后拼接的方式实现，但这种方法在分割图维度或类别数量变化时需要手动修改代码，所以想找一种通用的无循环实现方案：

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

无循环解决方案 #

利用NumPy的广播特性可以实现通用的无循环生成，无需手动指定每个类别，且自动适配类别数量和分割图维度：

方法1：自动识别所有非背景类别 #

import numpy as np

# 提取分割图中的所有唯一类别，并去除背景（这里背景为0）
classes = np.unique(a)[1:]

# 通过广播实现逐类别比较，直接得到(类别数, 1, 6, 6)的掩码数组
b = (a == classes[:, None, None, None]).astype(np.int32)

方法2：手动指定目标类别 #

如果不需要自动识别，而是要指定特定类别，只需替换classes数组即可：

# 指定需要生成掩码的类别
classes = np.array([15, 16, 13, 9])
b = (a == classes[:, None, None, None]).astype(np.int32)

原理说明 #

• classes[:, None, None, None]将类别数组的维度从(N,)扩展为(N, 1, 1, 1)，与原分割图a的(1, 1, 6, 6)维度匹配，触发NumPy广播机制，实现逐类别与整个分割图的元素比较。
• 比较后的布尔数组通过astype(np.int32)转换为0/1的二值掩码，最终维度为(N, 1, 6, 6)，完全符合需求。

内容的提问来源于stack exchange，提问作者Qimin Chen