np.argmin返回错误索引值：分割掩码转类别标签异常排查

背景 #

使用Paris Cityscapes数据集，输入图像尺寸为(256,512,3)，左半部分是输入图像，右半部分为对应的分割掩码。核心逻辑是通过匹配像素与31个类别的RGB值的最接近程度，将像素映射到classID。

掩码转换代码 #

为实现分割掩码到类别标签图像的转换，编写了transformMask函数：

def transformMask(input):
    global x
    input = np.expand_dims(input, axis= 2) #x.shape -> 1,1,31,3 | input.shape -> 256,256,1,3 (after expand_dims))
    output = np.argmin(np.linalg.norm(x - input, axis= -1), axis= -1)
    if np.sum(output > 31) > 0:
        exit(0)
    return np.broadcast_to(np.expand_dims(output, axis= -1), (256,256,3))

其中x是包含31个类别RGB值的数组。函数调用方式：

target_mask =  transformMask(image[:,256:,:]).astype(dtype= np.uint8)

异常现象 #

理论上，np.linalg.norm(x - input, axis=-1)的维度应为(256,256,31)，经np.argmin处理后输出维度为(256,256)，索引值应在0-30之间。但实际转换后的掩码出现大于31的值，且单独处理单张图像时无此问题。

可能的原因分析 #

• 全局变量x被意外修改：x的预设维度是(1,1,31,3)，但批量处理时如果其他代码逻辑修改了x的shape（比如误改成非预期维度），会导致广播计算异常，argmin取到超出0-30范围的错误索引。单张处理时没有触发这个修改，所以正常。
• dtype转换的溢出问题：调用时最后将结果转为np.uint8，如果output在转换前因计算错误出现了大于255的数值，转成uint8会发生截断溢出，表现出大于31的异常值。批量处理时的维度错误更容易触发这种溢出。
• 广播机制的意外匹配：若x的维度不符合预期（比如变成(31,1,3)），和input的(256,256,1,3)做减法时，广播后的维度会混乱，导致np.argmin(axis=-1)的计算逻辑错误，取到错误的索引值。
• 批量数据中存在异常像素：部分批量图像的掩码里存在不属于预设31类的RGB值，当x的维度异常时，argmin无法正确匹配到有效类别索引，从而产生超出范围的值。

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