高质量目标分割模型的关键在于掩码标注准确、数据处理一致、训练流程可复现;需确保掩码为单通道uint8二值图像(0/255)、尺寸与原图严格对齐,采用扁平化数据结构,图像与掩码同步增强,并优先验证掩码质量。
训练一个高质量的目标分割模型,关键不在模型本身多复杂,而在于掩码标注是否准确、数据处理是否一致、训练流程是否可复现。下面直接讲清楚从原始图像到最终掩码预测的完整链路,聚焦实操中真正卡点的地方。
多数主流分割模型(如Mask R-CNN、Segment Anything、YOLOv8-seg)要求掩码是单通道 uint8 图像,像素值为 0(背景)或 255(目标),且必须与原图尺寸严格对齐。常见错误包括:
建议用 OpenCV 快速校验并修复:
import cv2
mask = cv2.imread('mask.png', cv2.IMREAD_GRAYSCALE)
mask = (mask > 0).astype('uint8') * 255 # 强制二值化
assert mask.shape == img.shape[:2], "尺寸不匹配"
推荐采用 COCO 或 Detectron2 兼容的扁平结构,避免按类别建子目录(易导致 DataLoader 漏类或打乱顺序):
001.jpg → 001.png)这样写 Dataset 类最干净,也方便后续加新样本——只要丢进对应文件夹+更新 txt 即可。
普通 torchvision.transforms 不支持 mask 联动。必须用支持语义掩码的库,比如 albumentations:
import albumentations as A
transform = A.Compose([
A.Resize(640, 640),
A.HorizontalFlip(p=0.5),
A.RandomBrightnessContrast(p=0.2),
], bbox_params=A.BboxParams(format='coco', label_fields=['category_ids']),
mask_params=A.MaskParams(format='full'))
# 应用时传入 image 和 mask 一起
augmented = transform(image=img, mask=mask)
img_aug, mask_aug = augmented['image'], augmented['mask']
注意:mask_params 必须显式声明,否则 mask 可能被插值模糊甚至错位。
训练初期不要急着看 mAP,先可视化几个 batch 的预测掩码叠加在原图上:
cv2.addWeighted 把 mask 转成半透明红色图层如果前 10 个 epoch 的预测掩码连大致轮廓都不对,大概率是数据路径错、标签没对齐、或 mask 值域不是 0/255——这时候调学习率没意义。
基本上就这些。掩码分割不神秘,核心就是“对得上、看得清、变一致”。把数据流理顺了,换模型只是改几行配置的事。