ARTICLE DETAIL

资讯详情

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

基于PyTorch与U-Net的遥感滑坡识别实践

基于PyTorch与U-Net的遥感滑坡识别实践 简介本资源是一套完整的遥感图像滑坡识别实战项目面向计算机、遥感、地信等专业本科生及研究生专为毕业设计、课程设计与期末大作业打造。项目基于PyTorch框架构建CNN分类模型含AlexNet、ResNet等主流结构实现对高分辨率遥感影像中滑坡目标的自动识别与分类代码经导师指导并获99分高分评价小白可直接运行调试。压缩包共10个文件7个Python源码、1个预编译pyc、1个训练好的.pth模型权重、1个数据配置json总大小51.6MB涵盖数据集划分split_data.py、模型定义、训练/验证/推理全流程脚本及标准化数据加载模块结构清晰、注释完整、依赖明确。目前已有211人下载学习配套数据集与预训练模型开箱即用显著降低遥感AI项目入门门槛特别适合缺乏实战经验但需快速交付高质量项目的同学。 我第一次认真做遥感滑坡识别是在一个灾后评估项目里。手里是一批分辨率为1米左右的航拍影像滑坡体分布在破碎的山体两侧露出来的裸土和坡脚堆积物混在一起靠人工在屏幕上一块一块勾图斑几百个样本勾下来眼睛确实顶不住。后来团队换了个思路直接用深度学习CNN做像素级分割影像喂进去模型输出一张跟原图一样大小的掩膜图每个像素标记为滑坡或非滑坡。这套流程跑通之后原来需要三五天的目视解译工作压缩到几个小时就能出一版初判结果。这篇文章就完整分享一下我基于PyTorch搭建的遥感图像滑坡识别项目包含源码结构、数据集处理、CNN网络选型、训练评估以及训练好的权重文件怎么直接拿去做推理。不管你是刚接触遥感、想用深度学习方法做地物提取的学生还是工程上需要快速落地滑坡检测能力的从业者这套方案都可以直接参考复现。1. 项目整体设计与技术选型思路1.1 滑坡识别本质上是个像素级二分类任务滑坡识别从任务形态上看不是一个“给整张图打个标签”的图像分类问题。遥感影像里一张图往往同时包含滑坡体、裸露地表、植被、水体、道路、建筑等多种地物真正需要提取的是滑坡区域的空间范围和边界轮廓。所以正确的问题定义应该是对影像中每一个像素判断它属于滑坡还是非滑坡。这类任务在CV领域叫语义分割Semantic Segmentation和自动驾驶里的道路分割、医学影像里的器官分割是同一个套路。为什么不能简单用目标检测框来解决检测框会引入大量非滑坡背景像素而且滑坡体大多呈不规则条带状分布边界凹凸不平一个矩形框会把相邻的裸岩、冲沟一起包进来后续要做体积估算或者灾害范围统计时误差非常大。分割模型则能给出像素级边界这样得到的滑坡面积、周长等几何参数才有实际业务价值。我见过不少从分类任务转到分割任务的人第一步就把问题定义搞拧了所以特意把它放在最前面说清楚。后面那一系列关于网络结构、损失函数、评估指标的选择全都建立在这个“像素级二分类”的定位上。1.2 为什么选CNN而不是更强的新架构当前CV领域Transformer类模型确实很火像Swin Transformer、ViT在很多benchmark上都超过了CNN遥感分割任务里也有不少人用SegFormer、TransUNet。但我在这个项目里最终选择了CNN准确说是U-Net作为骨干核心原因是数据量和算力不匹配。Transformer模型非常依赖大规模数据预训练而滑坡数据集往往只有几千甚至几百个标注样本勉强够一个小型CNN收敛扔给Transformer很容易过拟合。而且遥感影像切块之后单张尺寸通常在256到512像素局部纹理和空间连续性才是判读滑坡的关键线索CNN通过卷积核天然具备局部感受野和平移等变性在中小规模数据集上反而更稳。再看训练成本。同样一块GPUCNN一轮epoch消耗的时间和显存不到Transformer的一半调参试错的周期更短。对于滑坡识别这种标注样本稀缺、工程落地要快的场景先上CNN是非常实际的选择。如果后续数据量增长到几万张再考虑把骨干网络换成Transformer也不迟训练好的模型托管在权重文件里切换骨干只需要改模型构建代码和权重路径其他训练与推理管线完全不用动。1.3 为什么用PyTorch搭这套工程深度学习框架现在主流的就PyTorch和TensorFlow两家另外还有MindSpore、PaddlePaddle等国产框架。我在这个项目里选PyTorch理由很直接动态图机制调试体验对新手和研究员都友好代码里遇到shape不匹配或者NaN直接断点到出错位置就能看到中间变量不用像静态图框架那样把模型先“编译”一遍再执行。另一个关键点是PyTorch生态对遥感领域非常完善。torchvision能加载ImageNet预训练的backbone配合timm库可以几行代码更换骨干网络onnx导出、TorchScript部署的参考资料多遇到问题基本都能搜到解决方案。除此之外PyTorch在日常数据加载的灵活度上确实方便。遥感任务需要做大量自定义操作比如多波段影像读取、窗口切片、地理坐标信息保留这些在PyTorch的Dataset子类里可以随意重写而不用被框架自带的数据管线限制住。1.4 项目的核心产物与适用对象这个项目最终交付的是三类东西源码、数据集、训练好的模型文件。源码包括数据加载、模型定义、训练循环、评估脚本、推理脚本五个模块目录结构比较清晰。数据集经过清洗和标注划分好了训练集、验证集、测试集。训练好的模型是一个U-Net结构的权重文件加载之后可以直接对新的遥感影像做滑坡区域预测不需要再走一遍训练流程。适合参考这套方案的人有两类一类是遥感、地理信息相关专业的学生想学深度学习方法在地学场景的应用另一类是做地质灾害调查、应急响应的工程技术人员需要快速搭建一个能用的滑坡识别工具。如果你只有分类任务的经验也可以拿这个项目当语义分割的入门案例里面的损失函数和评估指标代码可以直接复用到其他分割场景。2. 数据集处理与样本构建2.1 源码里的数据组织方式数据集部分是这个项目最费人工的环节。原始素材是多时相高分影像部分来自公开数据源部分来自合作单位提供的航拍成果最终保留下来的样本大约覆盖了不同地形、不同光照条件、不同植被覆盖度的滑坡场景。标注采用LabelMe工具手动勾画滑坡体外轮廓导出成JSON格式后编写脚本转成PNG格式的掩膜图滑坡区域像素值为255背景为0。源码里的数据目录是这样一个结构data/ ├── train/ │ ├── images/ │ │ ├── 001.png │ │ └── ... │ └── masks/ │ ├── 001.png │ └── ... ├── val/ │ ├── images/ │ └── masks/ └── test/ ├── images/ └── masks/图片和掩膜文件名一一对应。训练时用的Dataset类会同时读取image和mask然后做相同的几何变换这一点是分割任务数据处理和分类任务最大的区别图像如果翻转了掩膜也必须跟着翻转否则模型学到的就是错位标签。源码里把所有几何变换都封装成同一个随机种子操作确保图像和掩膜始终同步。2.2 影像切片最容易踩坑的一步遥感影像通常非常大一张原始图可能是上万乘上万的像素不可能直接把整张图塞进GPU。所以需要先把影像切分成固定大小的图块再分块送入网络。切块尺寸我建议256或者512太小了上下文信息不够模型容易把一块单独的裸土误判成滑坡太大了占显存且训练速度慢512一般已经接近常规GPU的上限。切块还有一个容易忽略的点重叠。如果训练数据是严格无重叠地硬切滑坡体恰好被切成两半、边界像素严重截断模型很难学到完整的滑坡形态特征。我在生成训练样本时设置了64像素的重叠相当于每张图块多保留了周围一部分上下文推理阶段也采用同样的切片策略预测完再把结果拼回完整图重叠区域取平均概率而不是简单取最大值这样拼缝处不会出现明显的块状伪影。另外要特别注意训练集、验证集、测试集的切块不能来自同一空间位置。如果训练集和验证集是从同一张大图上切出来的它们只相距几十个像素验证结果会虚高实际部署到新影像上效果立刻打折扣。我在划分数据集之前先按大图的空间范围分开确保任何两张图块不存在跨集合的重叠。2.3 数据增强与类别不平衡处理滑坡样本天然稀缺靠手动标注能获得的样本量有限因此数据增强是提升模型泛化能力的必要手段。我用的增强包括水平翻转、垂直翻转、90度旋转任意角度、随机裁剪、亮度对比度微调、高斯噪声添加。其中随机亮度对比度增强特别关键因为不同时期的遥感影像光照条件差异很大模型见过足够多的亮度变化才能在新影像上保持稳定。滑坡识别还有一个绕不开的问题类别极度不平衡。整幅影像中滑坡像素往往只占百分之一甚至更少如果用普通交叉熵损失模型很快就会陷入“全部预测为背景”的局部最优因为这样损失也足够低。解决方案是在损失函数里加入Dice Loss或者给滑坡类分配更高的权重。我在源码里把BCEWithLogitsLoss和DiceLoss按1比1加权组合效果比单独用任何一种都好。Dice Loss直接优化预测区域与真实区域的重叠度对前景小目标更加敏感BCE则保证每个像素的梯度信号都存在防止Dice Loss在小目标上梯度不稳定。两个损失结合模型在前景稀疏时的收敛速度明显加快。3. 网络结构与核心代码拆解3.1 U-Net为什么适合这个任务U-Net原本是医学图像分割提出的结构但用在遥感地物提取上也非常合适。它由编码器和解码器两部分组成编码器连续做卷积和下采样逐层提取从低级边缘纹理到高级语义的特征解码器逐步上采样恢复分辨率把高层的语义信息映射回每个像素的分类结果。U-Net最重要的设计是跳跃连接。编码器每层下采样后的特征图会直接拼接到解码器对应层这样解码器在恢复细节时不仅能靠高层语义猜测位置还能参考编码器里保留的边缘纹理信息。对滑坡识别来说边缘细节决定边界精度跳跃连接能让模型把“这里是滑坡”的语义判断和“滑坡边界到这里结束”的纹理线索结合起来出图的边界质量比单纯从低分辨率特征上采样好很多。3.2 关键模块代码实现我用的U-Net输入是3波段RGB影像归一化到0到1范围。模型结构按经典实现从32个通道起步每层下采样后通道翻倍最深一层到256通道。下面给出核心的DoubleConv模块和整体前向流程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, 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.conv(x)class UNet(nn.Module): def __init__(self, in_channels3, out_channels1): super().__init__() self.enc1 DoubleConv(in_channels, 32) self.enc2 DoubleConv(32, 64) self.enc3 DoubleConv(64, 128) self.enc4 DoubleConv(128, 256) self.pool nn.MaxPool2d(2) self.dec4 DoubleConv(256 128, 128) self.dec3 DoubleConv(128 64, 64) self.dec2 DoubleConv(64 32, 32) self.outc nn.Conv2d(32, out_channels, 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)) d4 self.dec4(torch.cat([nn.functional.interpolate(e4, scale_factor2, modebilinear, align_cornersFalse), e3], dim1)) d3 self.dec3(torch.cat([nn.functional.interpolate(d4, scale_factor2, modebilinear, align_cornersFalse), e2], dim1)) d2 self.dec2(torch.cat([nn.functional.interpolate(d3, scale_factor2, modebilinear, align_cornersFalse), e1], dim1)) return self.outc(d2)这套结构看起来很简洁但实际上手写分割网络时最常犯的错是解码器跳跃连接时channel数量没对齐。代码里编码器输出的通道数与解码器上采样后的特征图通道数不一定相同所以用torch.cat拼接前一定要确认两侧dim1的通道数一致。上面代码在拼接前把解码器上一层的输出和对应编码器特征对齐了运行过程不会报错。3.3 损失函数与评估指标设计损失函数部分我用的是组合损失。普通的BCEWithLogitsLoss对每个像素独立计算损失在背景占比极高时失效DiceLoss则直接衡量预测掩膜和真实掩膜的重叠率。实际训练时两者按1比1相加既保留像素级梯度又让模型更关注前景区域。def dice_loss(inputs, targets, smooth1.0): inputs torch.sigmoid(inputs) inputs inputs.contiguous().view(-1) targets targets.contiguous().view(-1) intersection (inputs * targets).sum() return 1 - (2.0 * intersection smooth) / (inputs.sum() targets.sum() smooth) def combined_loss(inputs, targets): bce nn.functional.binary_cross_entropy_with_logits(inputs, targets) dice dice_loss(inputs, targets) return bce dice评估指标用IoU交并比和F1分数最合适。准确率在这种类别极不平衡的场景下没有参考价值因为模型全部预测为背景也能得到98%以上的准确率看起来好像很厉害实际什么都没学到。IoU计算预测为正类且真实为正类的交集面积除以并集面积对前景区域敏感得多。4. 模型训练、评估与推理部署4.1 PyTorch环境准备与版本匹配项目代码基于PyTorch 2.x完成训练时用GPU加速。最推荐的方式是用conda创建独立环境避免把系统Python环境搞乱。conda create -n landslide python3.9 -y conda activate landslide pip install torch2.0.1 torchvision0.15.2 pip install numpy opencv-python pillow tqdm scikit-learn需要注意PyTorch、CUDA、显卡驱动三者的版本匹配。安装前可以先查看驱动支持的CUDA版本用nvidia-smi查看右上角。我的机器驱动支持CUDA 12.1所以安装了对应PyTorch版本。如果驱动版本较老可能得用CUDA 11.8甚至11.3对应的PyTorch版本否则torch.cuda.is_available()会返回False。不确定的时候直接去PyTorch官网首页用配置工具选择操作系统和CUDA版本生成的安装命令比自己在命令行猜版本靠谱得多。4.2 训练超参与训练过程细节训练阶段的超参数是我多次实验后确认比较稳的一组优化器Adam初始学习率1e-4batch size8图像尺寸512时比较吃显存如果OOM降到4训练轮数80轮配合早停机制学习率调度ReduceLROnPlateau验证集IoU连续5轮不提升就把学习率降低为原来的0.5早停验证集IoU连续15轮不提升则停止训练再保存最优模型训练时日志输出每一项损失和验证IoU验证集IoU最高的权重保存为best_model.pth。我建议训练时用tensorboard或者简单的日志文件记录曲线。损失下降趋势如果震荡剧烈可以先检查学习率是否过大如果训练损失下降但验证指标不动大概率是过拟合第一时间增加数据增强力度或者减小模型容量比盲目加训练轮数有效。4.3 加载训练好的模型做预测模型训练完成后实际使用只需要加载权重文件不需要重新训练。推理脚本的核心逻辑很简单读取影像、切块、送入模型、得到概率图、拼接、二值化、输出矢量范围。import torch from model import UNet model UNet(in_channels3, out_channels1) model.load_state_dict(torch.load(best_model.pth, map_locationcuda)) model.eval() with torch.no_grad(): pred_logit model(image_tensor) pred_prob torch.sigmoid(pred_logit)这里有个容易被忽略的细节模型训练时如果用了BatchNorm推理时必须把模型切到eval模式否则BatchNorm会继续用当前batch的统计量导致输出结果不稳定。很多人在推理时报错说效果变差结果就是忘了model.eval()。4.4 后处理与精度提升技巧模型输出的概率图在0到1之间默认取0.5作为阈值得到二值掩膜但滑坡边缘区域概率通常比较模糊。我在项目里做了一步形态学后处理先对二值图做开运算去掉孤立噪点再做闭运算填补滑坡内部的细小空洞。更进一步的技巧是连通域面积过滤。滑坡体在影像上通常有一定规模那些面积小于某个阈值的独立区域大概率是误检。我用OpenCV的connectedComponentsWithStats统计每个连通域面积过滤掉小于500像素的小区域误报率能下降不少。再配合地理信息工具将像素掩膜转成GeoJSON或Shapefile得到的不只是图片上的色块而是能落到地图上的滑坡范围可以进一步做面积量算和空间分析。另外还有一个提升边界精度的做法多尺度推理。预测时把输入影像分别按原始尺寸、0.75倍、0.5倍缩放各推理一次再把三个概率图缩放到原始尺寸取平均。这个方法会带来一些额外计算量但对边界模糊的滑坡区域通常能换来1到2个点的IoU提升效果值得。5. 常见问题与排查技巧实录5.1 环境与版本问题问题torch.cuda.is_available()返回False。排查顺序先看驱动有没有装好nvidia-smi能不能正常输出再看PyTorch版本对应的CUDA是否和驱动兼容最后检查是否安装成了CPU版PyTorch。常见情况是pip install torch时默认装了CPU版本用官方CUDA安装命令重新装一遍就能解决。问题训练时报错“CUDA out of memory”。显存不足优先减小batch size到4甚至2其次考虑降低切块尺寸从512降到384。如果两者都不想牺牲可以用梯度累积模拟较大batch size但训练时间会变长。混合精度训练AMP也能有效降低显存占用PyTorch自带torch.cuda.amp模块把前向和反向包装在autocast里就行。5.2 训练过程问题问题模型一直预测全背景IoU为0。绝大多数原因是类别不平衡。先看损失函数是不是普通BCE如果是换成BCE加DiceLoss的组合然后在训练日志里打印出前景像素占比确认每个batch里有没有包含滑坡样本。我遇到过因为切块窗口随机采样很多batch里完全不含滑坡导致梯度方向被大量背景样本带偏。问题训练损失很低验证集IoU也高但拿到新影像上效果崩了。这大概率是数据集划分时空间重叠导致的过拟合。确认训练集和验证集是否来自同一张大图的不同区域如果是重新按图划分数据。另一个可能是数据增强不足模型对新的光照或传感器类型不适应增加亮度扰动和模糊噪声能改善。问题模型对滑坡边界预测不准外扩或者内缩明显。多尺度推理和后处理膨胀腐蚀能缓解一部分但更根本的是让模型看到更多边界处的上下文。可以尝试把切片尺寸从256放大到384或512让模型有更大感受野来理解滑坡与周边地物的空间关系。5.3 预测效果问题问题推理时对整张大影像预测拼出来的结果有明显的块状纹理。原因是切片时没有设置重叠或者重叠区拼接时直接取了最大值。改成64或128像素重叠重叠区域取平均概率块状伪影基本消失。问题预测结果包含大量小面积误检区域。优先用连通域面积过滤再检查一下训练数据的负样本是否足够丰富。如果训练数据里没有类似颜色的裸土、河滩、道路路堑边坡模型很容易把这些区域误判为滑坡。这种情况下需要补充负样本重新训练比单纯调后处理阈值更有效。问题模型对灰度影像或非标准RGB影像表现差。遥感数据源可能是四波段、八波段甚至十六波段影像我自己通常做法是只取R、G、B三个波段送入模型或者对多波段做PCA降到三通道。训练和推理时保持相同波段选择否则模型分布完全不同。如果必须用到多光谱信息需要修改模型输入通道数并用对应波段的数据重新训练。最后再说一个实操中的体会。滑坡识别这类遥感分割任务很多人一开始把精力全放在调网络结构上真正跑起来才发现数据处理和推理细节决定了下限。切片重叠、数据集空间隔绝、训练推理时波段一致、推理时记得model.eval()、后处理过滤小连通域这些看起来不起眼的点每一个都把结果拉高一点加起来效果差异非常明显。项目里打包好的源码、数据集和预训练权重可以直接复用但切换到新的区域或新的传感器数据时千万别直接依赖老权重建议至少用一小批新样本做验证再决定是直接部署还是微调。遥感数据的时间变化太大上个月的模型这个月可能就不太好用了这是做地学深度学习不得不接受的现实。本文还有配套的精品资源点击获取
返回列表