简介这是一份基于Python与Unet模型实现的图像语义分割项目资源面向希望学习深度学习视觉任务的学生或开发者适用于毕业设计、课程设计、大作业以及工程实训等场景。项目覆盖数据准备、模型训练与语义分割预测的完整链路包含数据集生成脚本、训练脚本、预测脚本并提供训练好的Unet模型h5权重文件可直接加载验证帮助使用者跳过繁琐训练快速看到分割效果。资源压缩包共24个文件总大小约479MB以14个png图像文件为主同时含4个py脚本、1个pptx项目说明、1个xml辅助信息和1个h5模型文件目录内按训练、预测、数据等模块组织便于分块学习。目前已有204人浏览学习。读者可通过该资源理解Unet结构在图像语义分割中的应用思路也能借鉴其数据生成与训练组织方式迁移到自己的分割任务中。1. pythonUnet模型的图像语义分割给每个像素分类的入门首选半天就能看到效果针对pythonUnet模型的图像语义分割我最早是在一个航拍建筑轮廓提取项目里认真用起来的。当时对比过FCN和PSPNet最后落到Unet上理由很直接结构直观、显存占用不高、训练曲线容易理解代码在两三百行内能写完。这类任务解决的是像素级分类问题输入一张普通图片输出每个像素属于哪个类别比如道路、建筑、植被、水体。它适合手里有几千张标注图、想两三天内跑通思路的工程师和学生也适合用Unet作为新想法的最小复现基线。对新手来说完整跑通一个Unet分割项目能同时理解数据集格式、训练循环和评估指标三件事。接下来我会按环境、数据、模型、训练、排错、验证的顺序把整个链路拆开讲清楚。2. 环境与数据准备先把torch跑通再把标注图转成单通道索引2.1 环境配置torch版本和显存大小的取舍做Unet语义分割的常见做法是用conda建一个独立虚拟环境避免和日常开发环境互相污染。我一般会这样建conda create -n unet python3.10 -y conda activate unet pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install opencv-python pillow numpy tqdm逻辑说明python选3.10是因为目前torch对它的支持最稳配合vscode配置python环境时直接在解释器里选这个虚拟环境的路径就行。torch安装那行如果用的是CPU机器把最后的cu118参数去掉直接pip install torch也能装但分割训练用CPU基本上只能验证代码跑通训练一个VOC量级的数据集要按天计。装完之后用python -c import torch;print(torch.cuda.is_available())检查输出True再继续。显存约束是最先要确认的。输入patch为256×256、batch_size等于8时Unet在训练状态下大约需要4GB显存同等配置batch_size等于16需要7GB左右。所以8GB显存的卡是入门标配12GB以上可以放心调大batch和patch。如果手里只有笔记本集成显卡建议先走CPU模式跑通预处理和单张推理训练放到带NVIDIA显卡的机器上。2.2 选数据与搭目录训练自己的数据集前先搞懂mask格式图像语义分割的数据集由两张图组成一张是原始图片一张是和图片等大的mask标注图。mask里每个像素的数值代表类别编号。以VOC为例总共21类背景是0飞机是1自行车是2person是15。训练时模型只看两件事原图里这个像素长什么样以及mask里这个像素被标成几。这是整个任务最基本的逻辑。准备自己的数据集时目录结构我习惯按下面这样放原图和mask同名不同后缀后续写DataLoader几乎不用改data/ train/ image/0001.jpg mask/0001.png val/ image/0001.jpg mask/0001.pngimage放原始JPGmask放PNG。PNG是无损格式能保存单通道索引JPG带压缩噪声不适合做标注图。很多人在这里踩的第一个坑是直接从标注工具导出的mask是彩色RGB图等于用颜色给类别编码。网络拿到彩色mask之后会把三个通道当作三类特征去学最后loss降不下去。所以把彩色mask转成单通道索引图是训练前必须做的一步。2.3 掩码重编码脚本把VOC彩色mask转成单通道索引图VOC的mask是彩色调色板图每个像素的RGB颜色对应一个类别。下面这段脚本能把VOC的mask转成训练用的单通道索引图import numpy as np from PIL import Image import os def voc_mask_to_index(rgb_path, out_path): img Image.open(rgb_path).convert(P) # 转成调色板模式 arr np.array(img) # 像素值就是调色板索引 Image.fromarray(arr.astype(np.uint8)).save(out_path) # 遍历 VOC 的 SegmentationClass 目录全部转换 for name in os.listdir(VOCdevkit/VOC2012/SegmentationClass): if name.endswith(.png): voc_mask_to_index( fVOCdevkit/VOC2012/SegmentationClass/{name}, fdata/train/mask/{name} )逻辑说明先让PIL以调色板模式打开彩色mask此时numpy数组的像素值就是调色板索引直接保存就得到单通道索引图。参数说明convert(P)这一步是关键如果直接convert(RGB)再取某一通道得到的不是类别索引训练必然出问题。保存格式用PNG保证无损。如果你是用LabelMe标注的数据导出的本来就是JSON可以写一个解析JSON的脚本把多边形内部填充成类别索引过程与上述转换等价核心思想相同。2.4 标注数据验收肉眼扫一遍再开训数据转换完成后别急着训练。写一个快速脚本把原图和mask并排对比显示循环看一百张。我习惯重点检查三类问题mask是否为单通道、类别编号是否超出预设范围、原图和mask是否对得上。这一步看似费事但能避免最贵的错误训练一天后发现标签错位。一个常见的自检手段是用numpy统计mask像素值的分布。比如VOC是21个类别如果发现最大像素值远大于20说明某些区域被错误地填成了255这类问题会造成模型输出类别数对不上。所以我每次准备完新数据集都会跑一条命令python -c from PIL import Image; import numpy as np; anp.array(Image.open(data/train/mask/0001.png)); print(a.dtype, a.min(), a.max())看到类别集合在预期范围内再进入下一步。这一步虽然是肉眼工程但能挡掉80%的数据源头问题值得养成习惯。3. 搭建Unet并跑训练结构原理、损失函数与三个必调参数3.1 Unet结构核心对称编码器和跳跃连接为什么对分割友好Unet结构从名字到实现都直白左边是编码器逐层下采样特征通道翻倍空间尺寸减半右边是解码器逐层上采样把空间尺寸加倍左右之间有四条跳跃连接把编码器对应层的特征图直接拼到解码器同尺度的张量上。这个设计解决的是下采样丢细节的问题。分割任务最怕边缘糊跳跃连接可以把浅层的轮廓信息直接送到深层让输出在保持语义的同时保留锐利边界。如果不想手写全部结构常见做法是直接用开源实现。无论用谁的代码都要先检查最后输出层的类别数是否等于你的数据集类别数。以VOC为例网络输出通道必须等于21每个通道预测一个类别的概率图最终argmax得到索引。类别数不匹配是新手最容易忽略的问题。我早期手写过Unet核心就这一个卷积块import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.conv(x)参数说明padding1保证卷积不改变特征图尺寸biasFalse是因为后面紧跟BatchNormBatchNorm自带可学习的偏移项再设bias是冗余的。BatchNorm不能去掉否则深层特征和数据分布偏移会叠加模型很难收敛到好的mIoU。下采样时用maxpool上采样时用双线性插值加卷积具体组合见第4.3节。3.2 损失函数与训练循环一张训练脚本的参数明细图像语义分割默认的损失函数是交叉熵。交叉熵逐像素计算对每个像素的预测和真实类别做比较。但当类别不平衡时比如背景占90%、道路占10%模型会倾向把所有像素都预测成背景loss照样下降mIoU却不高。所以我在实际训练中会给交叉熵配置class_weight按每个类别在训练集中出现的频率倒数作权重具体计算方式在第4.4节展开。训练循环我的写法大致如下import torch import torch.nn.functional as F def train_one_epoch(model, loader, optimizer, device, class_weightNone): model.train() for images, masks in loader: images images.to(device) # [B, 3, H, W] masks masks.to(device) # [B, H, W]单通道索引 optimizer.zero_grad() logits model(images) # [B, C, H, W] loss F.cross_entropy( logits, masks, weightclass_weight, ignore_index255 ) loss.backward() optimizer.step() return loss.item()关键点masks形状是[B, H, W]不要传one-hot编码PyTorch的cross_entropy直接接受类别索引ignore_index255用来跳过边缘填充区域或未标注区域让模型忽略这些位置的预测而不是去猜类别。参数上我一般起点是Adam加lr1e-3batch_size8patch_size256训练20到30个epoch。学习率是最常翻车的点1e-3适中如果loss一开始就震荡立即降到1e-4。训练过程中每个epoch在验证集上算一次mIoU只保存验证mIoU最高的权重训练loss最低不等于泛化最好。3.3 用自己的数据训练Dataset类怎么写最稳自己标注的数据最常见的形式是image和mask同名但不同后缀。Dataset类的核心就两件事读原图、读mask然后同步预处理。import os import cv2 import torch from torch.utils.data import Dataset class SegDataset(Dataset): def __init__(self, image_dir, mask_dir, size(256, 256)): self.images sorted(os.listdir(image_dir)) self.image_dir image_dir self.mask_dir mask_dir self.size size def __len__(self): return len(self.images) def __getitem__(self, idx): name self.images[idx] img cv2.imread(os.path.join(self.image_dir, name)) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img cv2.resize(img, self.size, interpolationcv2.INTER_LINEAR) mask_path os.path.join(self.mask_dir, name.replace(.jpg, .png)) mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) mask cv2.resize(mask, self.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这段代码最需要强调的是两个resize插值方式的区别原图用线性插值没毛病mask必须用最近邻插值。如果用线性插值去缩放mask物体边缘会产生介于两个类别之间的数值比如0.7和0.3的混合值round之后就会捏造出一个并不存在的交界类别这会污染标注。另外mask读取用IMREAD_GRAYSCALE确保是单通道索引。img除以255.0归一化到[0,1]mask不用归一化保持原始类别索引这是语义分割和普通目标检测在数据流上最大的不同。3.4 判断收敛loss和mIoU互相印证才敢早停训练时不建议只盯loss曲线。loss下降只能说明模型在拟合训练集验证集mIoU才能说明分割效果。我的训练习惯是每5个epoch保存一次模型同时记录训练loss和验证mIoU。如果连续5个epoch验证mIoU没有上升就加载之前保存的最好权重把学习率降到原来的十分之一继续训练。关于早停还有一个细节把训练日志导出成CSV用pandas画两条曲线一条是训练loss一条是验证mIoU。当发现loss降到0.2左右不再动但mIoU还在40%到60%之间徘徊时大概率不是模型结构的问题而是小目标类别占比过低这就要回到损失函数和类别权重去调整而不是盲目加epoch。4. Unet语义分割的常见问题排查显存溢出、标签错位、棋盘格伪影这一章是Unet使用时的注意事项里最容易让人耗一整天的部分。以下四条都是我实际踩过的坑每条按现象、原因、解决的顺序写照着排查能省下大量时间。4.1 显存溢出先把patch_size降到256别硬扛现象训练刚开始没几步程序报CUDA out of memory代码通常停在backward附近。原因Unet的一次性前向传播会保留很多中间特征图patch_size512、batch_size16对12GB显存来说非常危险加上loss计算产生的中间变量显存峰值往往出现在backward时。新手第一反应是换更大的显卡其实多数情况是参数配置没有做减法。解决降低patch_size最立竿见影。从512降到256显存占用几乎降到四分之一。如果项目必须用原图大尺寸把batch_size降到2甚至1同时配合梯度累积。我常用的是bacth_size4、累积2次等效batch_size还是8但峰值显存小了一半。注意梯度累积时要把总loss除以累积步数否则等效学习率偏大容易导致训练震荡。4.2 mask还是彩色图训练标签必须是单通道索引现象训练时loss从一开始就在2.0以上波动训练完可视化预测结果时整张图颜色分布混乱看不出任何类别的连贯区域。原因数据准备环节没做掩码重编码直接把彩色RGB mask喂给了损失函数。交叉熵期望输入是[B, H, W]的单通道索引收到[B, H, W, 3]后某些框架会自动展平把三个通道当成三个样本维度网络输出类别数对不上整个梯度也就乱了。解决回到掩码转换脚本检查训练数据的mask到底是不是单通道。一条命令就能验证python -c from PIL import Image; import numpy as np; anp.array(Image.open(data/train/mask/0001.png)); print(a.shape)输出(H, W)就是对的(H, W, 3)就是错的。这个问题我踩过不止一次之后的习惯是每次换数据集第一次训练前先打印一个batch的mask.shape看到[B, H, W]才放心进训练循环。4.3 边缘出现棋盘格伪影解码器上采样方式要调整现象分割结果整体轮廓正确但物体边缘出现规律的格子状花纹像棋盘一样尤其在小目标密集区域。原因Unet解码器里如果直接用nn.ConvTranspose2d做上采样转置卷积在特征图重叠区域容易产生不均匀的权重分布棋盘格伪影就是这么来的。转置卷积不是不能用但需要额外的初始化技巧和精细调参对初学者性价比很低。解决把转置卷积替换成nn.Upsample(modebilinear)加普通卷积。我实现Unet时统一用这个组合先双线性上采样放大两倍再接一个3×3卷积平滑。效果比转置卷积稳定伪影基本消失。上采样后加卷积层会显著减少边缘的频谱噪声参数增加不多换来干净的边缘非常值得。4.4 loss在降但mIoU不涨检查类别不平衡和class_weight现象训练loss从2.5降到0.6验证mIoU始终停在20%上下前景目标一点都分割不出来。原因数据里背景占90%以上目标只占10%交叉熵被背景主导。模型学会了把全局预测为背景loss也很低但小目标类别全丢了。解决给损失函数加class_weight权重按频率倒数归一化。计算方式不复杂import numpy as np pixel_cnt np.zeros(num_classes) # 统计每个类别的像素总数 class_weight 1.0 / (pixel_cnt 1e-6) class_weight class_weight / class_weight.sum() * num_classes还有一个补充办法是改用混合损失0.5倍的交叉熵加0.5倍的Dice Loss。Dice Loss对前景占比小的类别更敏感能明显改善小目标场景。需要提醒的是加了class_weight之后loss的绝对值会变大这不代表模型变差判断模型好坏以mIoU曲线为准。5. 用单张推理与mIoU验证Unet效果一个趁手的验收技巧5.1 从模型权重到可视化推理阶段的最简实现训练完权重后我习惯先拿一张验证图做完整单图推理确认分割效果看得见再跑整个验证集的mIoU统计。原因是单张图能直观看到边缘是否光滑、小目标是否完整、类别是否被错分这些是数值指标描述不出来的。import cv2 import torch import numpy as np from PIL import Image def predict_and_overlay(model, img_path, palette, device): model.eval() img cv2.imread(img_path) img_rgb cv2.cvtColor(img, cv2.COLOR_BGR2RGB) input_tensor torch.from_numpy( img_rgb.transpose(2, 0, 1) ).float().unsqueeze(0) / 255.0 input_tensor input_tensor.to(device) with torch.no_grad(): logits model(input_tensor) pred torch.argmax(logits, dim1).squeeze(0).cpu().numpy() mask_color palette[pred.astype(np.uint8)] overlay cv2.addWeighted(img_rgb, 0.6, mask_color, 0.4, 0) Image.fromarray(overlay).save(overlay_result.png) return predpalette是一个[num_classes, 3]的numpy数组顺序和训练时的类别一一对应。addWeighted的alpha取0.6到0.7左右原图权重太高mask看不清太低边缘不直观。推理时默认按原尺寸输入如果训练时固定用了256×256这里要先resize推理后再把预测结果resize回原尺寸。验证集mIoU统计建议用numpy的bincount累加混淆矩阵逐像素for循环会慢到怀疑人生。我现在每训练一个分割模型都强制自己先跑一次单图推理看得见效果才敢跑全量验证这个习惯帮我避开了无数次“指标还行、图不能看”的尴尬。希望帮到你。本文还有配套的精品资源点击获取