使用SAM模型训练自定义数据集时遇掩码非二进制及ValueError问题

问题根源 #

1. 掩码读取时默认加载为3通道彩色图，resize操作引入渐变像素值，导致数值范围覆盖0-255而非预期的二进制（0/1）
2. 部分掩码为空或全0，触发训练逻辑中对空数组的min/max操作，抛出ValueError: zero-size array to reduction operation minimum which has no identity

具体修复步骤 #

1. 修正掩码读取与处理逻辑 #

修改掩码加载代码，确保读取单通道灰度图，合并后强制二值化：

import numpy as np

# 替换原Mask处理代码段
mask = None
for m in masks:
    # 读取单通道灰度掩码
    extr_mask = cv2.resize(cv2.imread(m, cv2.IMREAD_GRAYSCALE), (256, 256))
    # 二值化：所有非0像素设为1
    extr_mask = (extr_mask > 0).astype(np.uint8)
    if mask is None:
        mask = extr_mask
    else:
        # 合并多个掩码，保留所有前景区域
        mask = cv2.bitwise_or(mask, extr_mask)
# 兜底处理无掩码的情况，创建全0数组
if mask is None:
    mask = np.zeros((256, 256), dtype=np.uint8)
else:
    # 再次确认掩码为二进制格式
    mask = (mask > 0).astype(np.uint8)
mask_list.append(Image.fromarray(mask))

2. 优化数据集加载逻辑 #

无需将掩码列转为Image类型，直接存储numpy数组格式更适配SAM训练：

from datasets import Dataset

def create_dataset(images, masks, label_ids):
    dataset = Dataset.from_dict({
        "image": images,
        "mask": [np.array(m) for m in masks],
        "label": label_ids
    })
    # 仅将图像列转为Image类型
    dataset = dataset.cast_column("image", Image())
    return dataset

dataset = create_dataset(images=img_list, masks=mask_list, label_ids=label_list)

3. 添加掩码有效性校验 #

在训练前过滤或标记全0掩码样本，避免触发空数组操作错误：

def check_mask_validity(example):
    # 标记掩码是否包含前景区域（非0值）
    example["has_foreground"] = np.any(example["mask"] == 1)
    return example

dataset = dataset.map(check_mask_validity)
# 过滤无前景的样本（正常乳腺样本可根据需求选择保留或处理）
dataset = dataset.filter(lambda x: x["has_foreground"])

额外说明 #

• 二值化操作确保掩码严格只有0（背景）和1（前景）两个值，完全匹配SAM训练对ground truth的要求
• 若需保留全0掩码的正常样本，需在训练数据处理逻辑中添加判断，避免对空前景区域执行min/max等聚合操作

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