ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

宠物图像语义分割:数据集处理、U-Net训练与mIoU评估实战

宠物图像语义分割:数据集处理、U-Net训练与mIoU评估实战 简介面向图像分割入门与算法验证的大型宠物语义分割数据集适合计算机视觉学习者、算法工程师以及需要标准分割基准的科研场景。压缩包内共有两千个文件其中一千九百九十八个为PNG格式既包含原始图像也包含对应掩码另附一个类别说明文本文件和一个可视化脚本包体大小约七百六十六兆字节。数据已经预先做好训练集与验证集划分标注覆盖背景、宠物和边缘三类像素级标签通过类别文件可以核对每个类别的像素取值下载后几乎不需要额外整理便能直接接入Unet、Swin-Unet等常见分割模型开展训练与评估。随包提供的脚本会随机抽取一张图片将原始图像、真实标注掩码和标注叠加在原图上的蒙版效果并列展示帮助快速检查标签质量。目前已有四十八人学习适合用来熟悉分割数据组织方式、调试训练流程或作为课程实验与算法比对的基准数据。1. 宠物图像语义分割数据集为什么 7000 张图比一万张普通图更省事做图像分割的人最怕的不是模型选不对而是数据集到手先花两个礼拜洗数据。这套大型宠物图像语义分割数据集约7000张训练集和验证集已经按 imagesmasks 目录分好标签只有背景、宠物、边缘三类属于很典型的语义分割粒度。你拿到手不用再操心划分、重命名、类别映射直接接 U-Net 或 SwinUnet 就能跑通第一个迭代。适合刚入门分割的读者、要复现论文时快速验证网络结构的人以及想用一个干净数据把训练管线跑顺的从业者。比起类别多、标注复杂的遥感图像语义分割数据集三类的宠物数据更容易定位问题出在模型还是数据上。2. 数据集结构与标签编码先读懂 mask 像素值再想训练2.1 目录层级图片、mask、classes 文件三大块这类数据集的目录结构一般不需要再整理下载后解压直接能用。用 tree 命令铺开通常是这样的pet_seg_dataset/ ├── train/ │ ├── images/ │ │ ├── Abyssinian_32.png │ │ ├── Abyssinian_37.png │ │ ├── Bengal_1.png │ │ ├── Bombay_140.png │ │ ├── British_Shorthair_275.png │ │ └── Egyptian_Mau_120.png │ └── masks/ │ ├── Abyssinian_32.png │ ├── Abyssinian_37.png │ ├── Bengal_1.png │ ├── Bombay_140.png │ ├── British_Shorthair_275.png │ └── Egyptian_Mau_120.png ├── val/ │ ├── images/ │ └── masks/ └── classes.txttrain 下大概 5100 张图片和对应 maskval 下大概 2200 张。文件名是品种加编号比如 Egyptian_Mau_182.png、Abyssinian_172.png。因为图片和 mask 放在两个平铺目录里配对关系完全靠文件名所以训练前必须检查两边数量一致且同名文件一一对应。常见做法是只用 basename 排序不要相信压缩包解压后的默认顺序。类别定义不要靠猜classes.txt 里写明了像素值对应关系。以三分类为例通常是背景为 0、宠物为 1、边缘为 2但不同项目可能调整顺序所以我会先 cat 一下cat classes.txt head -n 3 classes.txthead 只看前三行就够先确认 0、1、2 各自代表什么。很多训练脚本会把 num_classes3 写死如果 classes.txt 里还出现第四个值说明这份数据不是纯三分类要立刻停下来。如果想确认目录数量可以在 bash 下这样查echo train images: $(ls train/images | wc -l) echo train masks: $(ls train/masks | wc -l) echo val images: $(ls val/images | wc -l) echo val masks: $(ls val/masks | wc -l)这样能直接看到 5100/5100、2200/2200 是否对得上。wc -l 按换行计数文件名里尽量不要有空格如果文件名带空格建议先批量改名否则后面写 DataLoader 时容易把路径截断。2.2 三类标签的像素约定与边缘类的价值三类标签的对应关系可以整理成下面这张表类别常见像素值语义在训练中的作用背景0图像中不属于宠物的区域提供负样本让模型学会“哪里不是宠物”宠物1猫或狗的主体区域分割目标决定区域召回率边缘2宠物轮廓或边界带约束边缘提升掩膜边界质量背景和宠物是常规语义分割类别。边缘单独拎出来是这套数据的特色等于把轮廓信息直接显式交给模型而不是让 U-Net 自己从损失函数里硬学。实际训练时边缘的像素占比通常只有几个百分点类别不平衡问题比二分类突出很多。很多人在第一版代码里用默认 CrossEntropyLoss结果背景学得不错宠物轮廓却糊成一团就是因为边缘类没有获得足够梯度。如果把三类理解成三通道语义通道0是背景概率通道1是宠物通道2是边缘。后续可视化、mIoU 计算都要围绕这个顺序组织不要在中途改变通道含义。一旦模型输出和 classes.txt 对不上所有评估都是白算。2.3 用脚本快速核对 mask 通道与像素分布写一个独立检查脚本是开始训练前最值的十分钟。我一般会读五张 mask确认它们是单通道灰度图且像素值只有 0、1、2import cv2 import numpy as np from pathlib import Path mask_dir Path(train/masks) mask_files sorted(mask_dir.glob(*.png))[:5] for mask_path in mask_files: mask cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE) print(mask_path.name, mask.shape, np.unique(mask))常见输出是(256, 256) [0 1 2]或者(512, 512) [0 1 2]。如果出现[ 0 255]说明这个 mask 是二值图不是三分类标签如果 shape 是(H, W, 3)说明被保存成了三通道 PNG要改成单通道再训练。单通道 mask 才能直接作为 CrossEntropyLoss 的 target三通道标签会让 loss 在类别维度上算错。再进一步统计全部训练集的类别比例判断是否值得给边缘加权class_count np.zeros(3, dtypenp.int64) for mask_path in sorted(mask_dir.glob(*.png)): mask cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE) for cls in range(3): class_count[cls] int((mask cls).sum()) print(class_count / class_count.sum())class_count 按 0、1、2 三类累计如果背景占了 95% 以上损失函数里不给边缘加权的话网络很容易把边缘类忽略掉。这个统计结果后面设权重时直接当依据。3. 数据读取与 Dataset 实现把 image 和 mask 对齐是第一道坎3.1 文件名对齐不要直接 zip 两个 listdir我见过最典型的翻车是把两个目录的 os.listdir 结果直接 zip 在一起然后训练很久才发现猫脸长在狗身上。原因是两个目录虽然文件内容同名但 listdir 返回顺序不一定一致尤其在 Windows 与 Linux 间拷贝过之后。正确做法是显式按文件名排序并加断言。import os image_dir train/images mask_dir train/masks image_names sorted(os.listdir(image_dir)) mask_names sorted(os.listdir(mask_dir)) assert len(image_names) len(mask_names), 图像和mask数量不一致 for img_name, mask_name in zip(image_names, mask_names): img_stem os.path.splitext(img_name)[0] mask_stem os.path.splitext(mask_name)[0] assert img_stem mask_stem, f{img_name} 与 {mask_name} 不配对这个脚本读起来像防守型编程但分割数据集里它真的能拦住很多问题。前面 sorted 保证了两个目录都按同一套字符串规则排序后面断言能快速暴露错位不用等训练到一半才看到 loss 异常。文件名里如果带下划线或数字sort 按字符串比较不会有问题只要图片文件名和 mask 文件名本身一一对应这套检查就成立。我还会顺手检查图片尺寸和 mask 尺寸是否一致。有些数据集原图是 224×224mask 却是 320×320Dataset 里直接 resize 也能强行配对但会引入肉眼看不见的对齐偏差。建议加一个最小检查import cv2 img cv2.imread(os.path.join(image_dir, image_names[0])) mask cv2.imread(os.path.join(mask_dir, mask_names[0]), cv2.IMREAD_GRAYSCALE) print(img.shape, mask.shape)如果宽高不一致后续统一 resize 到同一个尺寸时长宽比会被拉变形。常见做法是先 center crop 到正方形再 resize或者干脆统一按短边缩放到目标尺寸。对宠物图像这种主体居中的数据直接 resize 也能接受但需要知道这个前提。3.2 自定义 Dataset 类的最小实现下面的类可以直接抄到训练工程里。它做了三件事建立文件名映射、读取并缩放图像和 mask、把 mask 转成 long tensor 供交叉熵使用。import os import cv2 import numpy as np import torch from torch.utils.data import Dataset class PetSegDataset(Dataset): def __init__(self, image_dir, mask_dir, image_size(256, 256)): super().__init__() self.image_dir image_dir self.mask_dir mask_dir self.image_size image_size self.image_names sorted(os.listdir(image_dir)) self.mask_names sorted(os.listdir(mask_dir)) assert len(self.image_names) len(self.mask_names) for img_name, mask_name in zip(self.image_names, self.mask_names): assert os.path.splitext(img_name)[0] os.path.splitext(mask_name)[0], \ f{img_name} 与 {mask_name} 不对应 def __len__(self): return len(self.image_names) def __getitem__(self, idx): img_name self.image_names[idx] mask_name self.mask_names[idx] img cv2.imread(os.path.join(self.image_dir, img_name)) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) mask cv2.imread(os.path.join(self.mask_dir, mask_name), cv2.IMREAD_GRAYSCALE) img cv2.resize(img, self.image_size, interpolationcv2.INTER_LINEAR) mask cv2.resize(mask, self.image_size, interpolationcv2.INTER_NEAREST) img torch.from_numpy(img.transpose(2, 0, 1)).float() / 255.0 mask torch.from_numpy(mask).long() return img, mask逻辑说明集中在三个地方。第一mask 必须用cv2.IMREAD_GRAYSCALE读成单通道不能像读图片一样默认三通道。第二cv2.resize对 mask 用cv2.INTER_NEAREST对图片用线性插值最近邻不会在原来像素值 0、1、2 之间产生小数比如把边缘 2 和背景 0 的边界插值成 1这是分割标签里很隐蔽的错误来源。第三mask 返回时转成long()因为 PyTorch 的CrossEntropyLoss要求 target 是整型且形状为(B, H, W)而不是 one-hot。参数说明image_size(256, 256)是我在消费级显卡上常用的值内存占用小一个 batch 16 也能跑。如果原图比例不是正方形resize 会直接拉伸宠物主体一般居中影响不大如果后续要做精细边缘评估建议改成(512, 512)代价是显存占用翻两倍。num_classes3没有在 Dataset 里写死是因为后面损失函数和评估函数都会用到拆开写更清晰。3.3 训练/验证 DataLoader 与数据增强的同步问题数据集已经分好了 train 和 val不要自己在 train 里再切一个验证集出来。直接按目录构造即可from torch.utils.data import DataLoader train_dataset PetSegDataset(train/images, train/masks, image_size(256, 256)) val_dataset PetSegDataset(val/images, val/masks, image_size(256, 256)) train_loader DataLoader(train_dataset, batch_size16, shuffleTrue, num_workers4) val_loader DataLoader(val_dataset, batch_size16, shuffleFalse, num_workers4)shuffle 只开在训练集验证集固定顺序这样每次评估结果可复现。num_workers 按 CPU 核数调Windows 上建议小于或等于 4否则有时会卡在 DataLoader 初始化。另一个常见问题是增强。做随机翻转、随机缩放时图片和 mask 必须使用同一套变换参数。很多人用 torchvision 的独立 transform 分别处理 image 和 mask翻转方向就不一致了。更稳的做法是对 img 和 mask 用同一组参数import random if random.random() 0.5: img cv2.flip(img, 1) # 水平翻转 mask cv2.flip(mask, 1)这是最不容易写错的增强方式。水平翻转不会改变像素类别只需要让两个数组一起翻转。垂直翻转也可以但对宠物照片来说水平翻转更符合真实分布颜色增强只对 img 做不要动 mask。4. 损失、评估与可视化边缘类决定分割质量的天花板4.1 损失函数给边缘类设置更高权重宠物分割三分类里背景通常占大面积边缘占极小面积。用默认 CrossEntropyLoss 时网络每次迭代从背景学到大量“不是宠物”的梯度边缘像素对 loss 的贡献被稀释。最简单的缓解是给类别权重import torch import torch.nn as nn device torch.device(cuda if torch.cuda.is_available() else cpu) class_weights torch.tensor([0.5, 1.0, 2.0], devicedevice) criterion nn.CrossEntropyLoss(weightclass_weights)权重顺序和 classes.txt 保持一致背景 0.5、宠物 1.0、边缘 2.0。边缘给到 2 不是拍脑袋是因为从 2.3 节的像素统计里观察到边缘占比可能只有几个百分点给 2 到 5 之间都算常见范围。权重越大模型越倾向于把像素预测成边缘如果设得过高会出现整张图边缘过粗。如果后续用带权 Dice Loss 与交叉熵混合CrossEntropyLoss 的权重可以适当降低比如背景 0.3、宠物 0.8、边缘 1.5避免两种损失都强拉边缘造成过拟合。具体数值要靠验证集 mIoU 调不存在通用最优解。4.2 评估指标mIoU 比准确率更能暴露问题三分类里背景占比可能超过 80%所以整体准确率很容易冲到 90% 以上但宠物区域可能完全没分割出来。mIoU 是更可靠的指标。逐类 IoU 的计算如下import numpy as np import torch def iou_per_class(pred, target, num_classes3): ious [] for cls in range(num_classes): pred_mask (pred cls) target_mask (target cls) inter (pred_mask target_mask).sum().item() union (pred_mask | target_mask).sum().item() ious.append(inter / union if union 0 else float(nan)) return ious # pred 是 argmax 后的标签图target 是 GT pred torch.argmax(logits, dim1).cpu().numpy() target mask.cpu().numpy() ious iou_per_class(pred, target) miou np.nanmean(ious)逻辑说明IoU 把预测区域和真实区域的交集除以并集类别不平衡不会影响单个类别的计算只是背景 IoU 通常高、边缘 IoU 通常低差异正好能定位问题。union0的类在测试集里没出现返回 nan 再 nanmean避免把它算成 0 拉低整体。更严谨的做法要在整个验证集上累计而不是一张图一张图地平均。每张图只含一只宠物小图上的低 IoU 会被平均稀释大图边缘占比更敏感。我一般会把所有验证集 batch 的 pred 和 target 收集起来最后统一调用iou_per_class这样 mIoU 与逐类 IoU 更接近论文里的报告口径。4.3 可视化脚本原图、GT、叠加图三张一起存资源包里已经带了一个可视化脚本随机抽一张图把原始图片、GT、GT 蒙板在原图上的叠加结果保存到当前目录。这个思路很好但很多人的 GT 图保存后是全黑的因为 0、1、2 直接按灰度显示基本不可见。复制下面的映射逻辑import matplotlib.pyplot as plt import numpy as np idx 17 img, mask train_dataset[idx] img_np img.permute(1, 2, 0).numpy() mask_np mask.numpy() mask_rgb np.zeros((mask_np.shape[0], mask_np.shape[1], 3), dtypenp.uint8) mask_rgb[mask_np 1] [255, 200, 0] # 宠物主体 mask_rgb[mask_np 2] [255, 0, 0] # 边缘 overlay img_np.copy() overlay[mask_np 0] img_np[mask_np 0] * 0.6 mask_rgb[mask_np 0] / 255.0 * 0.4 fig, axes plt.subplots(1, 3, figsize(15, 5)) axes[0].imshow(img_np) axes[0].set_title(Original) axes[1].imshow(mask_rgb) axes[1].set_title(GT) axes[2].imshow(overlay) axes[2].set_title(Overlay) plt.savefig(visual_check.png, dpi150)逻辑说明mask 的像素值 1 和 2 在灰度显示下分别对应很暗的灰肉眼几乎看不出边界映射成黄色和红色后宠物轮廓一眼可见。overlay 部分只对非背景像素做融合背景保持原图这样能看到宠物位置是否和真实场景对齐。alpha0.6/0.4是我常用的叠加比例边缘如果需要更醒目可以把 0.4 提到 0.6。保存文件时建议固定idx而不是每次随机避免复现问题时要翻好几张图。资源自带脚本用随机抽取适合快速扫一遍自己做回归验证时应改成固定 idx让每次输出可比较。4.4 一个最小训练循环骨架把上面三块串起来最小可运行循环如下适合先确认数据管线没有错model UNet(in_channels3, num_classes3).to(device) optimizer torch.optim.Adam(model.parameters(), lr1e-4) for epoch in range(10): model.train() total_loss 0.0 for images, masks in train_loader: images images.to(device) masks masks.to(device) optimizer.zero_grad() logits model(images) loss criterion(logits, masks) loss.backward() optimizer.step() total_loss loss.item() print(fepoch {epoch}: loss {total_loss / len(train_loader):.4f})这里的 UNet 指你手头已经实现的网络比如资源配套参考里提到的 U-Net/SwinUnet/TransUNet 改进版本直接替换成你的模型类即可。逻辑说明mask 已经是(B, H, W)的 long tensor不需要 one-hotlogits 是(B, 3, H, W)cross entropy 内部会按类别维计算。参数说明lr 1e-4 是分割任务的保守起点如果训练 5 个 epoch 后 loss 完全不动把 lr 调到 3e-4 再试。5. 训练翻车避坑五个高发问题与排查顺序这一章的排查顺序是我按出现概率排的。遇到异常先按 5.1 到 5.5 的顺序检查而不是一上来就怀疑模型结构。大多数问题出在标签类型、文件配对、DataLoader 返回类型这三层。5.1 Loss 下降但 mIoU 始终为 0标签被归一化成了小数现象训练 loss 正常从 1.0 降到 0.4但验证 mIoU 打印出来是 0可视化预测 mask 全黑。原因最常见的是在__getitem__里对 mask 也执行了/ 255.0标签从 0、1、2 变成 0、0.0039、0.0078然后转成 long 时全部截断成 0等于所有像素都是背景。解决mask 保持原始整数像素值只对图像做归一化。在 Dataset 里检查torch.unique(mask)是否等于 0、1、2如果只有 0先回看预处理里有没有除 255。5.2 图像和 mask 错位宠物轮廓对不上现象训练曲线正常但可视化里猫的轮廓跑到了旁边另一只狗身上。原因image_names 和 mask_names 用了不同的排序规则或两个目录的文件列表经过拷贝后顺序不一致zip 时一对多错位。解决统一用sorted(os.listdir())并在 Dataset 构造时逐对断言os.path.splitext(img_name)[0] os.path.splitext(mask_name)[0]。这个断言失败时会立刻告诉你哪一对文件名对不上不用等训练完再猜。5.3 训练报错 expected scalar type Long 或 target shape 不匹配现象PyTorch 在计算 loss 时报expected scalar type Long或者报Expected target size [...]。原因mask 被读成三通道 RGB 后直接返回形状变成(3, H, W)而不是(H, W)或者 mask 是 float 类型CrossEntropyLoss 不接受。解决读取时用cv2.IMREAD_GRAYSCALE返回前加torch.from_numpy(mask).long()。如果用的是 PILImage.open(mask_path).convert(L)也可以但不要convert(RGB)。5.4 验证集指标虚高宠物个体在训练集和验证集重复现象训练 mIoU 只有 0.75验证 mIoU 却 0.95明显不合理。原因拿到全量数据后自己用 random_split 重新划分同一只宠物的不同照片同时进了训练和验证集。宠物个体相似度极高模型见过这只猫后验证集里它另一张照片几乎等于背答案。解决直接使用资源包已有的 train/val 目录。如果非要重新划分至少先按文件名里的品种编号去重保证同一 ID 只出现在一个集合里。5.5 GT 可视化全黑0、1、2 在灰度图里不可见现象用 matplotlib 直接 imshow mask显示出来是一块黑屏只有零星的灰点。原因0、1、2 在灰度 colormap 的默认范围 0-255 里分别对应黑、近黑、很暗肉眼几乎分辨不出宠物轮廓。这不是标签坏了只是显示映射问题。解决可视化前先把像素值重映射到 0、127、255或者做成 RGB 伪彩色。训练时仍然用原始 0、1、2不要因为显示问题去改标签文件。6. 批量预测与掩膜导出用边缘类 IoU 验证模型是否真的学会6.1 推理输出argmax 之后才是标签图训练完做推理时模型的 logits 形状是(1, 3, H, W)3 对应背景、宠物、边缘。需要沿类别维取最大索引model.eval() with torch.no_grad(): logits model(img.unsqueeze(0).to(device)) pred torch.argmax(logits, dim1).squeeze(0).cpu().numpy()argmax(dim1)的意思是在每个像素点上比较背景、宠物、边缘三类分数谁分数高就把像素归给谁。得到的 pred 是(H, W)的 int 数组值域 0、1、2。6.2 批量保存结果并统计输出占比一张张保存预测掩膜时建议直接存成单通道 PNG方便后续用专业工具检查import cv2 pred_vis np.zeros(pred.shape, dtypenp.uint8) pred_vis[pred 1] 200 pred_vis[pred 2] 255 cv2.imwrite(pred.png, pred_vis)如果把边缘类映射成 255肉眼就能看到轮廓。保存前要确认这套映射只用于展示不需要回传训练。想快速判断一个 batch 是否普遍预测失败可以统计每张图里三类像素占比unique, counts np.unique(pred, return_countsTrue) ratio {int(c): float(cnt / pred.size) for c, cnt in zip(unique, counts)} print(ratio)如果某张图 ratio 里只有背景没有宠物和边缘大概率推理图输入错了或者模型退化成了全背景分类器。6.3 看指标时不要只看 mIoU先看边缘类 IoU这是我在这套数据上摔过最狠的一跤。整段训练里验证集总 mIoU 到 0.82看起来不错后来把同一套流程挪到一个地毯图像分割的实时系统上轮廓一塌糊涂。回去翻逐类 IoU发现边缘类只有 0.17宠物类 0.86背景 0.94。mIoU 被背景和宠物拉高了边缘的低分完全被掩盖。正确的验收顺序是先打印三类 IoU再从三个数字里挑最低的看。对这套宠物数据边缘类 IoU 至少要过 0.3模型才有实用价值低于 0.2 说明边缘基本没学到。这套数据集的价值不在于让你今天跑出一个 90% 的 mIoU而在于用最少的变量把分割流程验证透。下载后先跑一遍可视化脚本确认三类标签显示正常再用验证集把 0/1/2 三类 IoU 都打出来最低的一类才是你真正要填的坑。从那以后我每次训练完都强制跑一遍逐类 IoU最低类别不过线就不存档。希望帮到你。本文还有配套的精品资源点击获取
返回列表