简介：这是一份基于Python与Unet模型实现的图像语义分割项目资源，面向希望学习深度学习视觉任务的学生或开发者，适用于毕业设计、课程设计、大作业以及工程实训等场景。项目覆盖数据准备、模型训练与语义分割预测的完整链路，包含数据集生成脚本、训练脚本、预测脚本，并提供训练好的Unet模型h5权重文件，可直接加载验证，帮助使用者跳过繁琐训练快速看到分割效果。资源压缩包共24个文件，总大小约479MB，以14个png图像文件为主，同时含4个py脚本、1个pptx项目说明、1个xml辅助信息和1个h5模型文件，目录内按训练、预测、数据等模块组织，便于分块学习。目前已有204人浏览学习。读者可通过该资源理解Unet结构在图像语义分割中的应用思路，也能借鉴其数据生成与训练组织方式，迁移到自己的分割任务中。
1. python+Unet模型的图像语义分割：给每个像素分类的入门首选，半天就能看到效果
针对python+Unet模型的图像语义分割，我最早是在一个航拍建筑轮廓提取项目里认真用起来的。当时对比过FCN和PSPNet，最后落到Unet上，理由很直接：结构直观、显存占用不高、训练曲线容易理解，代码在两三百行内能写完。这类任务解决的是像素级分类问题，输入一张普通图片，输出每个像素属于哪个类别，比如道路、建筑、植被、水体。它适合手里有几千张标注图、想两三天内跑通思路的工程师和学生，也适合用Unet作为新想法的最小复现基线。对新手来说，完整跑通一个Unet分割项目，能同时理解数据集格式、训练循环和评估指标三件事。接下来我会按环境、数据、模型、训练、排错、验证的顺序把整个链路拆开讲清楚。
2. 环境与数据准备：先把torch跑通，再把标注图转成单通道索引
2.1 环境配置：torch版本和显存大小的取舍
做Unet语义分割的常见做法是用conda建一个独立虚拟环境，避免和日常开发环境互相污染。我一般会这样建：
conda create -n unet python=3.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，自行车是2，person是15。训练时模型只看两件事：原图里这个像素长什么样，以及mask里这个像素被标成几。这是整个任务最基本的逻辑。
准备自己的数据集时，目录结构我习惯按下面这样放，原图和mask同名不同后缀，后续写DataLoader几乎不用改：
data/ train/ image/0001.jpg mask/0001.png val/ image/0001.jpg mask/0001.pngimage放原始JPG，mask放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( f"VOCdevkit/VOC2012/SegmentationClass/{name}", f"data/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; a=np.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, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True) ) def forward(self, x): return self.conv(x)参数说明：padding=1保证卷积不改变特征图尺寸；bias=False是因为后面紧跟BatchNorm，BatchNorm自带可学习的偏移项，再设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_weight=None): 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, weight=class_weight, ignore_index=255 ) loss.backward() optimizer.step() return loss.item()关键点：masks形状是[B, H, W]，不要传one-hot编码，PyTorch的cross_entropy直接接受类别索引；ignore_index=255用来跳过边缘填充区域或未标注区域，让模型忽略这些位置的预测而不是去猜类别。参数上，我一般起点是Adam加lr=1e-3，batch_size=8，patch_size=256，训练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, interpolation=cv2.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, interpolation=cv2.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_size=512、batch_size=16对12GB显存来说非常危险，加上loss计算产生的中间变量，显存峰值往往出现在backward时。新手第一反应是换更大的显卡，其实多数情况是参数配置没有做减法。
解决：降低patch_size最立竿见影。从512降到256，显存占用几乎降到四分之一。如果项目必须用原图大尺寸，把batch_size降到2甚至1，同时配合梯度累积。我常用的是bacth_size=4、累积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; a=np.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(mode='bilinear')加普通卷积。我实现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, dim=1).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循环会慢到怀疑人生。
我现在每训练一个分割模型，都强制自己先跑一次单图推理，看得见效果才敢跑全量验证，这个习惯帮我避开了无数次“指标还行、图不能看”的尴尬。希望帮到你。
本文还有配套的精品资源，点击获取
