简介:面向毕业设计场景的基于UNet的遥感图像语义分割完整实现,适合计算机视觉方向的本科生及研究生参考。资源内含69个文件,包含Python源码与pyc编译文件、Jupyter Notebook演示脚本、PNG示例图、LaTeX论文源文件以及PDF文档,包体约47.26MB,涵盖数据准备、模型构建、训练评估与推理预测的完整流程。已有346人学习下载。读者可获得可直接运行的UNet分割工程,包括create_dataset、train、predict等Notebook,以及模型定义、数据加载、训练代码和实验结果图;同时附有中文学位论文LaTeX源码与成品PDF,便于对照撰写毕业设计文档。整体目录结构清晰,兼具代码实践与论文写作参考价值。
1. 拿到遥感影像却分不出地物:UNet语义分割为什么是毕设和工程的首选
一张 0.5 米分辨率的遥感影像,尺寸动辄两三万像素见方,里面有耕地、建筑、水体、道路,你要的不是“检测出几个目标框”,而是把每一个像素都贴上一个类别标签。这就是语义分割,而 UNet 这种编码器-解码器结构,几乎是做这件事最稳妥的起点。不管你是为了完成毕业设计,还是想快速验证一个“遥感地物识别”的想法,UNet 都是先跑通流程、再谈改进的最短路径。
它胜在两点:一是结构简单,训练收敛快,个人电脑也能跑;二是跳跃连接让边界恢复得比普通全卷积网络好,最适合遥感里“道路窄、房屋边角锐利”这类细节。后面我会把从数据准备、网络复现到训练避坑的完整流程讲一遍,代码能直接抄,参数能直接套,最后告诉你训练完到底看哪些指标才算数。
2. 遥感图像和普通照片分割的三个差异:尺寸、波段和标注方式
2.1 影像尺寸:整图输入不可能,切片是躲不开的第一步
自然图像分割,比如 Cityscapes 数据集,图片通常是一两百万像素,缩放到 512×512 或 1024×1024 再进网络,内存开销完全可控。遥感影像不是这样,一景高分影像动辄 2 万×2 万像素,单通道就是 4 亿个像素点,直接缩放到 512×512 等于把一栋楼缩成一个点,地物细节全没了。所以业内常规做法是滑窗切片,把大影像切成 512×512 或 256×256 的小块,再送进网络。
切片有几个关键参数要提前定下来。patch size 建议 512×512,太小了感受野不够,建筑和道路这种中等尺度地物容易判错,太大了显存扛不住;overlap(重叠率)建议 0.25,也就是滑动步长为 384 像素,这样做预测时边缘不会出现明显的拼接缝。训练时可以用小一点的 overlap,甚至不重叠,但预测时必须有重叠,原因是图像边缘的像素上下文信息不完整,预测置信度低,重叠后取平均能明显改善边界处的碎斑。
2.2 波段数量和标注价值:RGB 之外还有四波段甚至更多
自然照片是三个波段(RGB),遥感影像常见的是四波段:R、G、B、近红外。对植被来说,近红外波段极其重要,归一化植被指数(NDVI)就是靠它算出来的,水体在近红外波段也几乎全吸收,和陆地区别更大。如果你的数据源给了四波段,不要只取 RGB 丢了近红外,把四个波段都用上,通道数从 3 变成 4,UNet 第一层卷积输入通道改一下就行,其他结构完全不用动。
如果你手上的数据是 16 位深度的 TIFF(单波段 0~65535),而预训练权重是按 8 位(0~255)设计的,直接喂进去会出现“均值偏移”的问题。我一般会先按百分比截断拉伸,比如把 2% 和 98% 分位数的像素值映射到 0 和 255,而不是全图最小最大值拉伸,否则影像里个别高亮的云或人造目标会把整个直方图拉垮,暗部地物全部挤在一起。
2.3 标注文件:别指望下载即用,多数标注是 GeoJSON 矢量
做语义分割,你需要的是和影像一一对应的逐像素类别标签图。公开数据集还好,模型训练标签通常是 PNG;但如果你用的是自己拿 QGIS 标的数据,或者某些数据商交付的标注,大概率是 GeoJSON 或 Shapefile 矢量面,需要栅格化才能变成 UNet 能吃的标签图。
用 GDAL 做矢量转栅格是常见方案,核心逻辑是:遍历每一个地物面,把它的类别属性值写到对应像素位置。类别 ID 从 0 开始,0 留作背景或“未分类”,类别从 1 开始编号。这个规则要和后面计算损失时的类别权重对齐,不然训练时模型会学出一个奇怪的偏置。
3. 用 PyTorch 从零搭一个 UNet:结构拆解与核心代码
3.1 DoubleConv 与编码器:下采样为什么用 max pooling 而不是 stride 卷积
UNet 编码器的基本单元是“两次卷积 + ReLU”,俗称 DoubleConv。我在实际代码里习惯给每一层加 BatchNorm,遥感影像不同切片的亮度差异很大(云雾遮挡、太阳高度角变化),BatchNorm 能让中间特征分布在训练中保持稳定,训练收敛速度明显更快,这也是新手最容易忽略的地方。不加热身的话,batch size 取 8 都容易训练不动。
下采样我保留经典设计的 max pooling,而不是换成 stride=2 卷积。max pooling 不引入额外参数,而且对纹理的微小位移更鲁棒。编码器每下采样一次,特征图尺寸减半、通道数翻倍,从 64 一路涨到 512。在显存不足时,优先砍第一层的通道数,比如从 64 降到 32,而不是减小输入 patch 到 256,因为 patch size 太小会牺牲上下文。
import torch 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), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, 3, padding=1), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), ) def forward(self, x): return self.conv(x)这段代码的要点:两次卷积都保持尺寸不变(padding=1),通道数由 in_ch 变成 out_ch。BatchNorm 在激活之前,epoch 数不多时效果比“卷积-ReLU-dropout”稳定得多,遥感数据 BN 的均值统计在推理时要用训练集的滑动平均,所以不要在推理前临时改训练模式。
3.2 解码器与跳跃连接:为什么说“边缘恢复就靠这一跳”
解码器每层先做一次转置卷积把特征图放大一倍,然后和编码器对应层做 channel 维度的拼接。编码器浅层特征保留了空间细节(边缘、纹理),深层特征更偏语义(这是建筑、这是水体),拼接后解码器同时获得“是什么”和“在哪里”两个信息。
我见过很多人在复现时抄错一个关键点:编码器的输出要保存下来,供解码器拼接用,而不是简单地把最终编码特征一路送到底。UNet 之所以在遥感分割上比 FCN 效果好,靠的就是这些跨层连接,把它们去掉,小目标(车辆、独立树木)的预测结果会明显残缺。
class UNet(nn.Module): def __init__(self, in_ch=3, num_classes=2): super().__init__() self.enc1 = DoubleConv(in_ch, 64) self.enc2 = DoubleConv(64, 128) self.enc3 = DoubleConv(128, 256) self.enc4 = DoubleConv(256, 512) self.pool = nn.MaxPool2d(2) self.bottleneck = DoubleConv(512, 1024) self.up4 = nn.ConvTranspose2d(1024, 512, 2, stride=2) self.dec4 = DoubleConv(1024, 512) self.up3 = nn.ConvTranspose2d(512, 256, 2, stride=2) self.dec3 = DoubleConv(512, 256) self.up2 = nn.ConvTranspose2d(256, 128, 2, stride=2) self.dec2 = DoubleConv(256, 128) self.up1 = nn.ConvTranspose2d(128, 64, 2, stride=2) self.dec1 = DoubleConv(128, 64) self.out = nn.Conv2d(64, num_classes, 1) def forward(self, x): e1 = self.enc1(x) e2 = self.enc2(self.pool(e1)) e3 = self.enc3(self.pool(e2)) e4 = self.enc4(self.pool(e3)) b = self.bottleneck(self.pool(e4)) d4 = self.dec4(torch.cat([self.up4(b), e4], dim=1)) d3 = self.dec3(torch.cat([self.up3(d4), e3], dim=1)) d2 = self.dec2(torch.cat([self.up2(d3), e2], dim=1)) d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1)) return self.out(d1)几个参数说明:
- in_ch 是输入通道数,RGB 就是 3,四波段就是 4,多光谱可以改成 8 甚至更多。
- num_classes 是类别总数,包含背景。比如你要分“耕地、建筑、水体”三类,再加上背景,这里写 4。
- 转置卷积 kernel 用 2、stride 用 2,输出尺寸正好翻倍,不会出现奇奇怪怪的尺寸对不齐。
- 最后一层是 1×1 卷积,把 64 通道压缩到类别数,后面接 CrossEntropyLoss 时不需要手动 softmax,PyTorch 的损失函数内部已经包含了 log_softmax。
3.3 损失函数:为什么交叉熵不够,DiceLoss 才是遥感的好搭档
遥感语义分割的类别分布极不均匀。一景影像里可能 70% 是农田,10% 是水体,3% 是道路,道路只有几个像素宽。如果直接用交叉熵,模型会倾向于把所有像素预测成农田,因为这样它的准确率也能到 0.7。这就是为什么光看 accuracy 会被严重误导。
我一般把 CrossEntropyLoss 和 DiceLoss 按 1:1 加权。DiceLoss 对小类别的惩罚更重,它能直接优化“预测区域和真实区域的重合度”,即使某一类的像素只占全图的 3%,它在损失里也能占到相当的比重。
class DiceLoss(nn.Module): def __init__(self, smooth=1.0): super().__init__() self.smooth = smooth def forward(self, logits, targets): probs = torch.softmax(logits, dim=1) num = targets.shape[0] # batch size # one-hot 编码 targets_onehot = torch.eye(logits.shape[1])[targets].permute(0, 3, 1, 2).to(logits.device) # 按 batch 和类别维度计算 dice intersection = (probs * targets_onehot).sum(dim=(2, 3)) union = probs.sum(dim=(2, 3)) + targets_onehot.sum(dim=(2, 3)) dice = (2.0 * intersection + self.smooth) / (union + self.smooth) return 1 - dice.mean()注意 targets 的 shape 是 (B, H, W),里面每个值是类别索引 0、1、2…;logits 的 shape 是 (B, C, H, W)。先对 targets 做 one-hot,才能和概率图逐像素相乘。smooth 设为 1.0 是为了防止分母为零,同时对类别极不平衡的情况有一定稳定作用。训练时 total_loss = 0.5 * ce_loss + 0.5 * dice_loss,两个损失的数值量级要观察一下,DiceLoss 的初始值通常在 0.3 到 0.7 之间,交叉熵初始值可能偏大,如果出现 CE 压过 Dice 的情况,把权重改为 0.4 / 0.6 再试。
4. 用自己的遥感数据集跑通 UNet:切片、标注对齐与训练参数
4.1 数据准备:把大影像切成 patch 时最容易犯的错
很多同学的毕设数据是自己标注的,或者在公开数据集上下载的“原图 + 标签图”组合。无论来源如何,第一步都是切片。我之前写过一版切片脚本,坑很多,这里直接给你一套稳定逻辑:先读取大图和标签图,然后按照 patch size 和 stride 遍历,切出来的小图和标签图都存入数组,再统一保存。注意两点:影像和标签图的坐标系必须一致(很多公开数据集的标签图是从矢量转的,像素尺寸和原图不完全对齐);切片后必须检查有没有“全背景”的 patch,这类样本要滤掉,否则训练时模型始终在看空白。
可以用滑动窗口遍历,同时在窗口内对标签图做一次判断:如果某个 patch 的标签图里前景像素占比低于 5%,直接丢掉。这样能显著提升训练效率,避免大量无效计算。
import numpy as np from PIL import Image def crop_patches(image, label, patch_size=512, stride=384): h, w = image.shape[:2] patches, labels = [], [] for y in range(0, h - patch_size + 1, stride): for x in range(0, w - patch_size + 1, stride): img_patch = image[y:y+patch_size, x:x+patch_size] lbl_patch = label[y:y+patch_size, x:x+patch_size] # 过滤背景占比过大的 patch if (lbl_patch == 0).sum() / (patch_size * patch_size) > 0.95: continue patches.append(img_patch) labels.append(lbl_patch) return np.stack(patches), np.stack(labels)这个实现是最直白的双层循环,慢但对小数据集够用。stride=384 意味着相邻 patch 有 128 像素重叠,训练数据量增大了约 30%,但模型在整图拼接预测时边界更稳。如果你用 stride=512,那就完全没有重叠,数据量小,训练速度快,但预测时的拼接缝要靠后处理去补,我建议训练时用 512、预测时用 384。
4.2 数据增强:专注几何变换,少做颜色抖动
遥感图像的语义分割,几何增强几乎是必需品。随机水平翻转、垂直翻转、旋转 90 度,这几个操作不改变类别语义,实现成本极低。我有一个习惯:旋转 90 度的倍数而不是任意角度,因为任意角度旋转会引入插值误差,标签图是离散的类别 ID,插值会把“1”变成“0.7”这种奇异值,标签图就毁了。
颜色增强要克制。适度做亮度扰动、对比度扰动可以模拟不同季节、不同太阳高度的成像差异,但不要做随机擦除或大幅度高斯噪声,遥感影像中地物边缘的噪点和自然照片不一样,随机擦除很容易把完整的建筑屋顶或者道路中间挖掉一块,模型会学到“不连续的地物”,这和你想要的效果正好相反。
4.3 训练参数:batch size、学习率、epoch 怎么定
我以 512×512 patch、单卡 12GB 显存为基准,给出一个可以直接上手的配置表。如果你的显存是 6GB,把 batch size 降到 2;如果是 24GB,试着把 batch size 提到 8,并适当增大学习率,收敛速度会明显提升。
| 参数 | 建议值 | 说明 |
|---|---|---|
| patch size | 512×512 | 太大显存溢出,太小物体破碎 |
| batch size | 4 | 12GB 显存下的稳妥值 |
| 初始学习率 | 1e-3 | Adam 优化器配合 |
| 学习率调度 | ReduceLROnPlateau | 验证 loss 连续 5 个 epoch 不降则乘 0.5 |
| epoch 数 | 60~100 | 遥感数据量小,100 轮一般足够 |
| optimizer | AdamW | 相比 Adam,weight decay 更规范 |
我见过很多人卡在“loss 一直降但验证 mIoU 不动”的怪圈,问题往往出在优化器和学习率。遥感数据集不像 ImageNet 那么大,初始学习率用 1e-3 是上限,超过这个值第一轮 loss 就可能冲到 NaN。如果用的是 SGD,学习率要降到 1e-2 附近,但收敛明显更慢,建议新手直接用 Adam,把 betas 设为默认值就好。
训练过程中每轮要记录两件事:训练集的 loss 和验证集的 mIoU。如果训练 loss 在降、验证 mIoU 也在升,说明一切正常;如果训练 loss 在降、验证 mIoU 停滞,说明过拟合,去增强数据或者降低模型容量(把第一层通道从 64 减到 48);如果训练 loss 不降,大概率是学习率设置不对,或者数据标签有错位。
5. 避坑:UNet 跑遥感数据时最常见的 5 个问题与排查
5.1 模型输出全黑或全背景:类别映射混乱
现象:训练完预测,输出图里几乎全是背景类,前景地物只偶尔出现几个亮点。
原因:最常见的是类别 ID 设置错误,比如你把标签值标成了 1、2、3,但模型输出通道是 3(不含背景),背景类 0 又没有像素参与监督;还有可能是标签图里 255 表示目标、0 表示背景,你直接把 255 当成一个类别喂给了网络,相当于所有目标都被当成同样的“第 255 类”,网络根本学不出区别。
解决:在切片之前先统计标签图里都有哪些像素值,print(np.unique(label)),一定要确认类别集合和损失函数里的类别数完全对齐。用一个统一的预处理函数把 255 重映射到 1,把 0 保持为 0,再检查一下切片后的标签 patch 里类别分布和整图一致。
5.2 训练 loss 正常但验证精度极低:训练集和验证集来自同一景影像的不同区域?
现象:训练集上 mIoU 已经到 0.75 了,验证集上只有 0.3,差距巨大。
原因:遥感影像的空间自相关性极强。如果你用同一景影像随机切 patch 后按 8:2 分成训练集和验证集,模型其实是在用“同一块地面”的上下文做记忆,验证集里很多 patch 和训练集距离非常近,甚至重叠。等拿到一景全新的影像上测试,性能立刻崩。
解决:验证集和训练集必须按“影像”划分,而不是按 patch 划分。同一景影像的 patch 要么全部进训练集,要么全部进验证集。最好使用不同日期、不同区域的数据做验证。交叉验证时也要按影像分 fold,不要按 patch 分。
5.3 预测图有密集的椒盐噪声和小碎斑:缺少后处理
现象:预测结果里单像素的类别点特别多,大块连续地物被很多杂点污染,像是被撒了盐。
原因:逐像素分类本质上对每个像素独立决策,没有相邻像素的约束。UNet 虽然通过卷积隐式利用了局部上下文,但它不会像条件随机场那样强制空间连续性,因此小碎斑在遥感里常见。
解决:快速方法是对概率图做高斯模糊,再重新取 argmax;更有效的是用形态学开闭运算来过滤碎斑。我通常的做法是:先对预测图做 mode filter(用 3×3 窗口取众数),再对少数类别(道路、水体)做闭运算连接断开的部分。不要把后处理做得太重,否则真实的小地物(独立树木、小汽车)也会被抹掉。
5.4 显存不够,batch size 降到 1 还是溢出
现象:一张 512×512 的 patch,batch size 设为 1,显存依然溢出。
原因:UNet 编码器第一层 64 通道、后续逐层翻倍到 1024,对显存的需求不低;但更常见的原因是反卷积层输出的中间特征图全部存在显存里用于反向传播,没有开启混合精度。
解决:加一行 torch.cuda.amp.autocast() 包住 forward 和 loss,用 GradScaler 做梯度缩放,显存占用直接下降约三分之一。另外检查一下是不是输入图片被读成了 float64,PyTorch 默认计算图用 float32,但如果你把图像转成 double,显存直接翻倍,这种情况常发生在把 16 位 TIFF 读进来没转 float32。
5.5 训练集和验证集的类别分布差异大:mIoU 虚高
现象:验证集的 mIoU 看起来很高,但逐类看,小类别的 IoU 只有 0.1 左右,大类却有 0.9。
原因:mIoU 是各类别 IoU 的算术平均,但如果验证集切片恰好切到了大面积的水体或农田区域,小类别在验证集里根本没出现几次,IoU 计算时的分母样本太少,波动极大。
解决:报告指标时至少要给出两个东西:每类 IoU(不平均)和加权 mIoU(按各类像素占比加权)。如果某一类在验证集里的像素总数不足 1000,这个 IoU 数值没有参考意义,单独标出来,别混进均值里。真正交付时,用一景完整影像做全图预测,再和人工标注算 IoU,这才是最终指标。
6. 训练完怎么判断模型能不能用:从混淆矩阵到交付一张带彩色标签的图
最后一关是精度评估和结果输出。mIoU 是一个粗糙的整体数,但你最好把混淆矩阵打印出来看一眼:道路被误判成什么?水体主要混到哪一类?如果道路大量被误判成阴影,但阴影不是你的目标类别,说明模型没有见过足够的道路阴影样本,需要补充标注,而不是调网络结构。
另一个容易漏的环节是类别 ID 到颜色的映射。训练时标签用的是 0、1、2,可视化时如果没有一份固定的 colormap,两张预测图的对比就没有意义。我会定义一个字典:0=黑色(背景)、1=绿色(植被)、2=红色(建筑)、3=蓝色(水体),然后输出成 PNG。同时输出一张半透明叠加在原始影像上的效果图,透明度 0.4,方便肉眼快速检查边界对得准不准。
import numpy as np from PIL import Image def save_prediction(prob_map, save_path, color_map): # prob_map: (H, W, C) 每个像素是各类别概率 pred = np.argmax(prob_map, axis=-1) # (H, W) rgb = np.zeros((pred.shape[0], pred.shape[1], 3), dtype=np.uint8) for cls_id, color in color_map.items(): mask = pred == cls_id rgb[mask] = color Image.fromarray(rgb).save(save_path)这不仅仅是出图好看,它帮你快速定位模型的系统性问题。如果建筑和道路在图上粘连不清,你就知道下采样的层数太深了,小物体细节在 bottleneck 里被压缩得太狠,可以试试把 UNet 的 depth 从 4 降到 3;如果是远处的细小道路断断续续,说明 patch size 偏小或者训练数据里这类样本太少,去增加“含道路”的 patch 数量比调网络更有效。
我自己的习惯是保存四个文件:预测标签 PNG、彩色可视化 PNG、半透明叠加图、以及每个类别的 IoU 数值表。这套东西齐全了,无论你是交给导师还是部署到一个小工具里,别人拿到都能秒懂结果好坏。我的血泪经验是:毕设评审看重的不是那个 mIoU 数字有多高,而是你能不能解释清楚每一类误差从哪里来,以及模型在整景影像上的泛化表现。如果你能从混淆矩阵里指出一条可操作的改进路线,这个项目基本就站稳了。希望帮到你。
本文还有配套的精品资源,点击获取