简介本资源为基于 Python 与 Unet 模型的图像语义分割完整项目包面向希望入门或进阶计算机视觉的开发者尤其适合作为毕业设计、课程设计、大作业或工程实训的参考方案。包内共 24 个文件以 14 张 png 图像样本、4 个 py 脚本、1 个 h5 训练权重、1 个 pptx 说明文档及 xml、ovr 等辅助文件为主压缩包约 478.98MB涵盖数据生成、模型训练、预测推理与结果可视化等环节。项目围绕 Unet 编码器—解码器结构展开提供从数据集构建到模型落地的完整链路读者可据此理解语义分割的数据标注、网络搭建、损失计算与精度评估思路并直接复用脚本进行二次开发。目前已有 204 人学习下载适合需要快速搭建分割实验环境、对照代码查漏补缺的初学者与进阶学习者。1. 从一张遥感图说起pythonUnet 图像语义分割到底在解决什么手里有一批无人机航拍图或者遥感影像想自动把建筑、道路、植被、水体一块块抠出来逐像素打上类别标签——这就是图像语义分割要干的事。传统做法靠阈值、纹理特征加分类器换个季节、换片区域就得重新调参泛化能力差得让人头疼。Unet 这套编码器-解码器结构之所以在分割圈子里经久不衰核心在于它的跳跃连接把浅层的高分辨率细节和深层的语义信息拼在一起小目标边缘不容易丢医学影像、遥感解译、工业缺陷检测都能用。python 生态里 PyTorch 或 TensorFlow 都能快速搭出 Unet配合 numpy、OpenCV 做数据预处理整条链路门槛并不高。这篇面向的是想自己动手跑通一套分割流程的工程师和新手从环境装起到数据组织、模型搭建、训练调参、推理可视化再到踩坑排查一步步给到能抄的代码和参数。遥感图像语义分割这类任务对显存和标注质量要求不低读完你应该能判断自己的数据值不值得上 Unet以及怎么少走弯路。2. 环境与数据把 pythonUnet 的输入管道先搭对2.1 python 环境与依赖安装的稳妥路径很多人卡在第一步不是模型是环境。python 安装教程网上一搜一大把但真到跑 Unet 的时候版本冲突才是血泪经验的来源。我一般用 conda 建独立环境避免和系统里其他 python 项目打架。python 3.8 到 3.10 对主流深度学习框架兼容性最好太新的版本有些库还没跟上。# 创建独立环境指定 python 3.9 conda create -n unet_seg python3.9 -y conda activate unet_seg # 安装 PyTorch按自己 CUDA 版本去官网查对应命令这里以 CUDA 11.8 为例 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 # 图像处理与科学计算常用库 pip install numpy opencv-python pillow matplotlib scikit-learn tqdm这段命令的逻辑是先隔离环境再装框架最后补图像处理库。参数上注意--index-url要和你机器上的 CUDA 驱动匹配装完用torch.cuda.is_available()验证返回 False 说明 CUDA 没接上别急着往下跑训练。numpy 是很多库的底层依赖python 安装 numpy 库的方法虽然简单但版本太新有时和 opencv 编译版本对不上遇到ImportError先查 numpy 版本。vscode python 环境配置或者 pycharm 都行关键是指定解释器到刚建的 conda 环境不然跑起来用的还是系统 python。2.2 语义分割数据集的组织与标注格式Unet 训练需要成对的「原图 掩码图」。掩码图里每个像素的值就是类别索引背景为 0目标从 1 开始编号。遥感图像语义分割常见做法是把大图裁成 256×256 或 512×512 的小块裁的时候要保证掩码和原图严格对齐否则训练 loss 降不下去还找不到原因。import os import cv2 import numpy as np def split_image_and_mask(img_path, mask_path, out_dir, patch_size256, stride128): 滑动窗口裁剪原图与掩码stride 小于 patch_size 可做重叠增强 img cv2.imread(img_path) mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) h, w img.shape[:2] idx 0 for y in range(0, h - patch_size 1, stride): for x in range(0, w - patch_size 1, stride): img_patch img[y:ypatch_size, x:xpatch_size] mask_patch mask[y:ypatch_size, x:xpatch_size] # 跳过全背景的块减少样本不均衡 if mask_patch.max() 0: continue cv2.imwrite(os.path.join(out_dir, img, f{idx}.png), img_patch) cv2.imwrite(os.path.join(out_dir, mask, f{idx}.png), mask_patch) idx 1 split_image_and_mask(raw/img.tif, raw/mask.png, dataset/train)裁剪逻辑里patch_size决定单样本视野太小语义信息不足太大显存吃不消256 是遥感任务的常用起点。stride设成 patch_size 的一半能增加样本量但重叠太多会让训练集和验证集泄漏划分时要以原图为单位切分不能按 patch 随机分。掩码必须是单通道灰度图类别值用整数别用调色板 PNG否则读出来是 RGB 三通道训练直接报维度错误。2.3 数据增强与 DataLoader 的落地写法分割任务标注成本高样本往往就几百张数据增强是刚需。但增强必须原图和掩码同步做几何变换翻转、旋转、缩放都要一致颜色抖动只作用在原图上。import torch from torch.utils.data import Dataset, DataLoader import albumentations as A from albumentations.pytorch import ToTensorV2 train_transform A.Compose([ A.HorizontalFlip(p0.5), A.VerticalFlip(p0.5), A.RandomRotate90(p0.5), A.RandomBrightnessContrast(p0.3), A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), ToTensorV2(), ]) class SegDataset(Dataset): def __init__(self, img_dir, mask_dir, transformNone): self.img_dir img_dir self.mask_dir mask_dir self.transform transform self.names os.listdir(img_dir) def __len__(self): return len(self.names) def __getitem__(self, i): name self.names[i] img cv2.cvtColor(cv2.imread(os.path.join(self.img_dir, name)), cv2.COLOR_BGR2RGB) mask cv2.imread(os.path.join(self.mask_dir, name), cv2.IMREAD_GRAYSCALE) if self.transform: aug self.transform(imageimg, maskmask) img, mask aug[image], aug[mask] return img, mask.long() loader DataLoader(SegDataset(dataset/train/img, dataset/train/mask, train_transform), batch_size8, shuffleTrue, num_workers4, pin_memoryTrue)albumentations 的 Compose 保证几何变换同步作用于 image 和 mask这是它比手写增强靠谱的地方。Normalize用的 ImageNet 均值方差如果从头训练可以换成自己数据集的统计值。batch_size受显存限制8 是 8G 显存的保守值num_workers在 Windows 上设 0 更稳Linux 可以开到 CPU 核数。mask 转long()是因为交叉熵损失要求标签是 int64。3. Unet 网络结构编码器、解码器与跳跃连接怎么写3.1 从 Unet 网络结构图看每一层的职责Unet 网络结构图那条经典的 U 形不是画着好看。左半边编码器每经过一次下采样特征图尺寸减半、通道翻倍感受野变大负责回答「这是什么」。右半边解码器上采样恢复分辨率负责回答「在哪里」。中间的跳跃连接把编码器同层特征直接拼到解码器对应层让边缘、纹理这些空间细节不被深层卷积抹掉。遥感图像里道路又细又长没有跳跃连接分割结果经常断成一段一段。import torch.nn as nn import torch.nn.functional as F class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.block nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.block(x) class Down(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.pool_conv nn.Sequential(nn.MaxPool2d(2), DoubleConv(in_ch, out_ch)) def forward(self, x): return self.pool_conv(x) class Up(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.up nn.ConvTranspose2d(in_ch, in_ch // 2, 2, stride2) self.conv DoubleConv(in_ch, out_ch) def forward(self, x, skip): x self.up(x) # 处理尺寸不整除导致的偏差 diff_y skip.size(2) - x.size(2) diff_x skip.size(3) - x.size(3) x F.pad(x, [diff_x // 2, diff_x - diff_x // 2, diff_y // 2, diff_y - diff_y // 2]) x torch.cat([skip, x], dim1) return self.conv(x)DoubleConv是 Unet 的基本砖块两次 3×3 卷积加 BN 和 ReLU。Down用最大池化下采样Up用转置卷积上采样后和跳跃特征拼接。拼接前那段 pad 是很多人忽略的坑输入尺寸不是 16 的整数倍时上采样后的尺寸和跳跃特征对不上torch.cat直接报错补零对齐是最省事的解法。3.2 完整 Unet 组装与输出层通道数的确定把砖块拼成完整网络输出层通道数等于类别数含背景。遥感图像语义分割如果分建筑、道路、植被、水体加背景就是 5 通道。class UNet(nn.Module): def __init__(self, in_ch3, n_classes5): super().__init__() self.inc DoubleConv(in_ch, 64) self.down1 Down(64, 128) self.down2 Down(128, 256) self.down3 Down(256, 512) self.down4 Down(512, 1024) self.up1 Up(1024, 512) self.up2 Up(512, 256) self.up3 Up(256, 128) self.up4 Up(128, 64) self.out nn.Conv2d(64, n_classes, 1) def forward(self, x): x1 self.inc(x) x2 self.down1(x1) x3 self.down2(x2) x4 self.down3(x3) x5 self.down4(x4) x self.up1(x5, x4) x self.up2(x, x3) x self.up3(x, x2) x self.up4(x, x1) return self.out(x)n_classes一定要和掩码里的最大类别值加一对应少一个通道训练时交叉熵会报 index 越界。in_ch按输入通道改RGB 是 3多光谱影像可能是 4 或更多。这个版本是最朴素的 Unetunet 模型改进通常从这里下手换编码器为 ResNet 或 EfficientNet 做迁移学习加注意力模块或者把转置卷积换成双线性插值上采样减少棋盘伪影。3.3 损失函数与评价指标的选型理由分割任务样本极不均衡背景像素往往占九成以上纯交叉熵会让模型学会全预测背景也能拿高准确率。常见做法是交叉熵加 Dice 损失组合。class DiceLoss(nn.Module): def __init__(self, smooth1e-6): super().__init__() self.smooth smooth def forward(self, logits, targets): probs F.softmax(logits, dim1) targets_onehot F.one_hot(targets, num_classeslogits.shape[1]).permute(0, 3, 1, 2).float() intersection (probs * targets_onehot).sum(dim(2, 3)) union probs.sum(dim(2, 3)) targets_onehot.sum(dim(2, 3)) dice (2 * intersection self.smooth) / (union self.smooth) return 1 - dice.mean() # 组合损失 ce nn.CrossEntropyLoss() def criterion(logits, targets): return ce(logits, targets) DiceLoss()(logits, targets)Dice 直接优化预测和真值的重叠度对小目标更友好。smooth防止分母为零。评价指标别只看像素准确率用 mIoU平均交并比更能反映真实分割质量每个类别算 IoU 再平均背景类通常不计入或单独看。4. 训练、推理与可视化把模型真正跑起来4.1 训练循环与学习率、batch size 的调参经验训练循环本身不复杂难的是参数怎么设。我一般用 Adam 起步学习率 1e-3 到 1e-4配合余弦退火或 ReduceLROnPlateau。import torch from tqdm import tqdm device torch.device(cuda if torch.cuda.is_available() else cpu) model UNet(in_ch3, n_classes5).to(device) optimizer torch.optim.Adam(model.parameters(), lr1e-3) scheduler torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, modemin, patience5, factor0.5) best_iou 0 for epoch in range(100): model.train() loop tqdm(loader, descfEpoch {epoch}) for imgs, masks in loop: imgs, masks imgs.to(device), masks.to(device) optimizer.zero_grad() logits model(imgs) loss criterion(logits, masks) loss.backward() optimizer.step() loop.set_postfix(lossloss.item()) # 每个 epoch 后在验证集算 mIoU这里省略验证函数细节 val_iou evaluate(model, val_loader, device) scheduler.step(-val_iou) if val_iou best_iou: best_iou val_iou torch.save(model.state_dict(), best_unet.pth)学习率太大 loss 震荡不降太小收敛慢到怀疑人生1e-3 是 Adam 的常用起点。patience5表示验证指标 5 个 epoch 不提升就砍半学习率。batch size 和显存直接挂钩显存不够优先降 batch 而不是降分辨率分辨率对分割精度影响更大。训练时盯着验证集 mIoU训练 loss 一直降但验证 mIoU 掉头向下就是过拟合加增强或加 dropout。4.2 推理脚本与分割结果可视化训练完要能对新图出结果推理时记得切 eval 模式和 no_grad。def predict(model, img_path, device, n_classes5): model.eval() img cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB) img cv2.resize(img, (256, 256)) tensor torch.from_numpy(img).permute(2, 0, 1).float().unsqueeze(0) / 255.0 tensor tensor.to(device) with torch.no_grad(): logits model(tensor) pred torch.argmax(logits, dim1).squeeze().cpu().numpy() return pred pred predict(model, test.png, device) plt.imshow(pred, cmapjet) plt.colorbar() plt.savefig(pred_vis.png, dpi150, bbox_inchestight)argmax取每个像素概率最大的类别得到分割图。可视化时 python 画图横坐标太密集是常见问题plt.xticks(rotation45)或者直接关掉坐标轴plt.axis(off)更适合看分割效果。类别多的时候用固定颜色映射表别用 jet颜色跳变容易误读。4.3 用 mIoU 和混淆矩阵验证分割质量光看图不够要量化。mIoU 按类别算混淆矩阵能看出哪两类容易混。def compute_iou(pred, target, n_classes): ious [] for cls in range(n_classes): pred_mask (pred cls) target_mask (target cls) inter (pred_mask target_mask).sum() union (pred_mask | target_mask).sum() if union 0: continue ious.append(inter / union) return np.mean(ious)union 0的类别说明验证集里根本没出现跳过不计入平均否则会被拉低。混淆矩阵用 sklearn 的confusion_matrix展平后算能直观看到建筑和道路是不是经常互相误判这类错误往往要靠增加难例样本或调整损失权重来解决。5. 避坑与排查Unet 训练里最容易翻车的几件事5.1 loss 不下降甚至变 NaN现象训练几个 batch 后 loss 变成 nan或者一直卡在高位不动。原因通常是学习率过大、数据没归一化、或者掩码里有超出类别数的值。解决先把学习率降到 1e-4 试检查输入像素是否归一化到 0-1 或标准化用np.unique(mask)确认掩码最大值小于 n_classes。混合精度训练时梯度溢出也会 NaN加torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)做梯度裁剪。5.2 验证集 mIoU 远低于训练集现象训练集 mIoU 0.9验证集只有 0.4。原因多半是过拟合或数据泄漏。解决先查数据划分重叠裁剪的 patch 如果随机分到训练和验证就会泄漏要以原图为单位划分。确认没泄漏后加数据增强、加权重衰减、或者换更轻的编码器。样本量太少时考虑迁移学习用预训练编码器初始化。5.3 分割结果边缘毛糙、小目标丢失现象大块区域分得还行细长道路断裂小建筑漏检。原因下采样太多次丢失空间细节或者损失函数对小目标不敏感。解决减少下采样层数或者用空洞卷积替代部分池化损失里加大 Dice 权重推理时用滑动窗口重叠预测再拼接边缘处取重叠区域的平均概率。5.4 显存溢出 OOM现象训练中途报 CUDA out of memory。原因batch size 太大、分辨率太高、或者验证时没加 no_grad 导致计算图累积。解决降 batch size用梯度累积模拟大 batch推理和验证包在torch.no_grad()里及时del中间变量并torch.cuda.empty_cache()。混合精度训练也能省不少显存。5.5 推理速度慢到无法落地现象单张图推理要好几秒。原因模型太大、没转半精度、或者用了 CPU。解决推理时model.half()转 FP16输入 tensor 也转 half导出 ONNX 或 TensorRT 加速确认torch.cuda.is_available()为 True别在 CPU 上硬跑。遥感图像语义分割如果要做大图推理分块并行比整图输入更实际。6. 进阶技巧让 Unet 在你自己数据上真正好用跑通基础版只是起点真正决定效果的是几个细节。第一编码器换预训练权重。从头训练 Unet 在小数据集上很难收敛好把编码器换成在 ImageNet 上预训练的 ResNet34 或 EfficientNet-B0解码器随机初始化mIoU 通常能涨 5 到 10 个点。第二损失函数按类别加权。遥感任务里道路、水体样本少给这些类更高的交叉熵权重或者用 Focal Loss 压制易分类样本。第三测试时增强TTA。推理时对同一张图做水平翻转、垂直翻转各预测一次概率图平均后再 argmax边缘稳定性明显提升代价是推理时间翻几倍看场景取舍。验证方法上别只信一个 mIoU 数字。我会把验证集里 IoU 最低的几张图单独拎出来看往往能发现标注错误或者某类目标系统性漏检。标注质量是分割任务的天花板模型再强也救不了错标。另外训练日志里除了 loss 和 mIoU把每个类别的 IoU 都打出来某一类突然掉下去通常对应数据里那类样本分布变了。一个具体技巧学习率预热。前 500 个 iteration 从 1e-6 线性升到设定值再正常衰减能明显减少训练初期的震荡尤其是 batch size 较大的时候。代码上用一个 lambda 函数控制 optimizer 的 lr 就行。def warmup_lr(step, warmup_steps500, base_lr1e-3): if step warmup_steps: return base_lr * step / warmup_steps return base_lr for step, (imgs, masks) in enumerate(loader): lr warmup_lr(step) for pg in optimizer.param_groups: pg[lr] lr # 正常训练步骤...我自己踩得最狠的一次坑是掩码用了调色板 PNG读出来是 RGB训练 loss 死活不降查了两天才发现是数据格式问题。从那以后我养成了一个习惯拿到任何分割数据集先np.unique看掩码值再cv2.imread确认通道数这两步花不了一分钟能省掉大量玄学调试。希望帮到你。本文还有配套的精品资源点击获取