Numpy判断某维度是否相同

在 语义分割（Semantic segmentation） 领域深度学习的label预处理工作中，要把3通道的label数据转为单通道的mask作为训练样本。
例如把一堆黑色和红色的像素，转为全是0和1的单通道 mask。

最Normal的处理流程为：

1. 判断像素的颜色；
2. 根据颜色赋值对应的label。

那么在步骤一中，就要判断 label 的 channel 维度（例如shape为(480, 1280, 3)，则channel维度为最后的3）和 color是否相同。

假设 label 大小为 4 x 4，3个颜色通道：

color_mask = np.random.randint(0, 10, (4, 4, 3))

print(color_mask)
"""
array([[[4, 1, 8],
        [7, 8, 0],
        [1, 8, 7],
        [9, 2, 0]],

       [[4, 8, 5],
        [7, 6, 6],
        [6, 5, 5],
        [2, 7, 4]],

       [[9, 2, 1],
        [3, 7, 6],
        [9, 4, 7],
        [5, 9, 1]],

       [[1, 9, 8],
        [6, 7, 9],
        [2, 9, 5],
        [3, 3, 5]]])
"""
print(color_mask.shape)
"""(4, 4, 3)"""

初始化 color ：

color = np.array([3, 3, 5])

print(color)
"""array([3, 3, 5])"""

基于 numpy 的矩阵特性，可以以矩阵为单位进行条件判断：

res = color_mask == color

print(res)
"""
array([[[False, False, False],
        [False, False, False],
        [False, False, False],
        [False, False, False]],

       [[False, False,  True],
        [False, False, False],
        [False, False,  True],
        [False, False, False]],

       [[False, False, False],
        [ True, False, False],
        [False, False, False],
        [False, False, False]],

       [[False, False, False],
        [False, False, False],
        [False, False,  True],
        [ True,  True,  True]]])
"""

观察到以上结果是在最后一维进行判断的。
使用np.all()，即可把 4x4x3 维度的布尔条件映射到 4x4 的结果中。
其中axis=n指定在第 n 维度上判断：

np.all(res, axis=-1)
"""
array([[False, False, False, False],
       [False, False, False, False],
       [False, False, False, False],
       [False, False, False,  True]])
"""

最后在预先初始化好的矩阵上赋值即可：

encode_mask = np.zeros(color_mask.shape[:2], dtype=np.long)
encode_mask[np.all(res, axis=-1)] = 1

print(encode_mask)
"""
array([[0, 0, 0, 0],
       [0, 0, 0, 0],
       [0, 0, 0, 0],
       [0, 0, 0, 1]])
"""

以上为使用 numpy 判断某维度是否相同的方法。