ARTICLE DETAIL

资讯详情

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

基于Pytorch的U-Net遥感滑坡识别项目实战

基于Pytorch的U-Net遥感滑坡识别项目实战 简介本资源是一套基于PyTorch框架实现遥感图像滑坡识别的完整深度学习项目面向地质灾害监测、遥感图像分析领域的研究人员与AI工程实践者解决滑坡区域自动定位与分类这一典型地物识别问题。压缩包共122个文件含18个Python源码涵盖CNN模型构建、RPN与Faster R-CNN核心模块、数据加载、训练与推理全流程、96个XML标注文件提供带边界框与类别标签的遥感样本、6个txt说明及配置文件、1个README.md文档和1个字体文件总大小4.93MB。已有58人下载学习。资源提供可直接运行的训练代码、预训练模型权重及标注完备的数据集配套详细项目说明文档清晰阐述模型设计逻辑、数据预处理规范、训练调参要点与评估指标解读便于快速复现、二次开发或教学演示。 从遥感解译转到地质灾害方向后我第一次拿到滑坡识别这个任务时心里其实挺没底的。遥感图像里滑坡的形态变化大加上植被、阴影、裸岩的干扰传统目视解译需要专业人员耗上大量时间去勾画边界。后来基于深度学习CNN网络用Pytorch框架做了这套滑坡识别项目整个流程跑通后发现效果和效率都比传统方法提升了一个量级。所以这篇就把项目的完整思路、数据怎么处理、模型怎么搭、训练时踩过哪些坑以及如何直接使用训练好的模型做推理一次性写清楚。无论你是刚开始接触遥感图像识别的学生还是需要快速落地滑坡检测的工程师这篇内容都能给你一个可以直接参考的路径。1. 项目概述与整体思路1.1 滑坡识别的实际需求滑坡灾害发生在山区、丘陵、河谷地带通常由降雨、地震或人类工程活动触发。在遥感影像上滑坡体往往表现为色调异常、纹理破碎、地形突变范围从几十平方米到几平方公里不等。传统解译方式依赖人工在影像上勾画多边形一个中等面积的滑坡可能需要半小时甚至更久而且解译结果受个人经验影响很大。时间紧、任务重的时候人工解译根本忙不过来。所以用深度学习模型自动识别滑坡区域本质上是把“看图说话”这件事交给网络让模型学习滑坡在光谱、纹理、形状上的特征然后输出像素级的分割结果。这个项目解决的就是“遥感影像上哪里有滑坡、边界在哪”的问题属于语义分割任务。1.2 技术选型为什么是CNN加Pytorch常见的遥感图像识别任务包括目标检测和语义分割。滑坡识别更需要得到完整的边界轮廓所以核心是语义分割。CNN在这里承担的是特征提取器的角色通过卷积、池化、激活等操作从原始像素中逐层提取低级的边缘纹理特征和高级的语义特征。Pytorch作为深度学习框架提供自动求导、GPU加速、动态计算图调试起来非常方便。我选择Pytorch的主要原因是它的生态足够成熟torchvision里有预训练模型可以直接加载社区资料也多遇到问题基本都能搜到解决方案。加上项目里需要频繁调整网络结构和训练策略Pytorch写起来比静态图框架更灵活更适合研究性质的工程落地。这个项目里我用了U-Net结构作为基础编码器部分换成ResNet34预训练权重。为什么不直接从头训练一个CNN因为遥感影像标注成本高数据量通常只有几千张从头训练容易过拟合而迁移学习能显著提升收敛速度和最终精度。1.3 拿到zip包后先看什么项目结构拆解项目压缩包解压后应该能看到这些核心内容landslide_cnn/ ├── data/ # 原始数据集与标注文件 │ ├── images/ # 遥感影像一般为tif或png │ ├── masks/ # 标签图二值图滑坡区域为白色 │ └── train_val_split.py # 数据集划分脚本 ├── models/ │ └── unet_resnet.py # 网络结构定义 ├── scripts/ │ ├── train.py # 训练脚本 │ ├── predict.py # 推理脚本 │ └── utils.py # 数据处理与评价指标工具 ├── weights/ │ └── best_model.pth # 训练好的模型权重 ├── requirements.txt └── 项目说明.pdf建议先打开项目说明确认数据集格式和模型输入尺寸再打开模型定义文件。我第一次跑这种项目时急着直接训练结果因为图像通道数和标签格式不匹配浪费了半天。遥感影像有可能是四通道R、G、B、近红外而模型预训练权重是在三通道RGB上训练的这一点需要特别留意。2. 数据准备与预处理2.1 遥感影像数据来源与标注规范数据处理是整个项目中耗时最长的环节。我用的数据集是某山区的高分辨率遥感影像分辨率为0.5米到2米不等。原始影像覆盖范围很大需要切分成小块才能送入网络训练。我采用的切窗大小是512×512像素因为U-Net下采样到1/16分辨率输入太小会丢失上下文信息太大则显存吃不消。标注规范方面滑坡区域统一用二值掩码表示滑坡体像素值为1背景为0。标签图片和原始影像必须保持相同尺寸和分辨率否则后续训练时矩阵对不上。常见的问题在于标注软件导出的形状是GeoJSON转成栅格时容易发生偏移建议每次转换后用叠加显示的方式检查一遍。如果你要自己标注新数据抠滑坡边界时尽量沿着滑坡后壁和堆积区的边缘走不要漏掉碎屑流的部分。可以用QGIS或者Labelme完成导出为多边形后再栅格化。2.2 数据增强与样本均衡滑坡样本在影像中通常只占一小部分正负样本比例可能高达1:10甚至更高。直接训练会让模型偏向预测背景所以要做两类处理一是数据增强二是损失函数加权。我用的增强方式包括随机翻转、旋转90度、随机亮度对比度调整、高斯噪声。这些操作可以在线进行也就是每个epoch动态生成不同的增强样本相当于无形中扩大了数据集。注意不要用随机裁剪因为裁剪会破坏滑坡的完整性。旋转角度我用90度的倍数避免引入插值噪声。另一个思路是拼接采样就是把包含滑坡的小块和纯背景块按照一定比例混合之后作为训练集。我在项目中按1:2的比例混合保证每个batch里至少有一半样本含有滑坡区域这样模型不会一开始就“躺平”。2.3 数据集划分与文件组织划分方式采用按区域划分而不是按图片随机划分。因为遥感影像相邻切片高度相关如果同一区域的切片同时出现在训练集和验证集验证结果会虚高模型实际的泛化能力会被高估。我按地理位置把整个研究区划分为四个大块三块做训练一块做验证。文件组织建议保持简单data/images/train/ data/images/val/ data/masks/train/ data/masks/val/脚本里写一个简单的数据集类读取图片和对应掩码训练时返回增强后的张量。要注意影像和掩码的文件名保持一致最好前缀相同比如img_001.tif对应mask_001.png这样脚本读取时不容易出错。3. 模型设计与训练实现3.1 网络结构设计从基础CNN到U-Net变体纯CNN分类网络只能输出图像级别的类别无法给出像素级滑坡位置。因此我选择U-Net这种经典的编码器-解码器结构。编码器不断下采样扩大感受野提取语义特征解码器通过上采样逐步恢复空间分辨率并与编码器的跳跃连接融合保留精细边界信息。滑坡的边界通常不规则所以跳跃连接很关键。模型定义的核心代码如下import torch import torch.nn as nn import torchvision.models as models class ResNet34UNet(nn.Module): def __init__(self, num_classes1, pretrainedTrue): super().__init__() resnet models.resnet34(pretrainedpretrained) self.encoder1 nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu) self.encoder2 nn.Sequential(resnet.maxpool, resnet.layer1) self.encoder3 resnet.layer2 self.encoder4 resnet.layer3 self.encoder5 resnet.layer4 self.center nn.Sequential( nn.Conv2d(512, 512, kernel_size3, padding1), nn.BatchNorm2d(512), nn.ReLU(inplaceTrue) ) self.decoder5 DecoderBlock(512, 256) self.decoder4 DecoderBlock(256, 128) self.decoder3 DecoderBlock(128, 64) self.decoder2 DecoderBlock(64, 64) self.decoder1 nn.Sequential( nn.Conv2d(64, 64, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.Conv2d(64, num_classes, kernel_size1) ) def forward(self, x): e1 self.encoder1(x) e2 self.encoder2(e1) e3 self.encoder3(e2) e4 self.encoder4(e3) e5 self.encoder5(e4) center self.center(e5) d5 self.decoder5(center, e4) d4 self.decoder4(d5, e3) d3 self.decoder3(d4, e2) d2 self.decoder2(d3, e1) out self.decoder1(d2) return torch.sigmoid(out)DecoderBlock里做上采样和跳跃连接拼接我用的双线性插值上采样而不是转置卷积因为转置卷积容易产生棋盘格伪影。3.2 损失函数与评价指标滑坡样本占比低单纯用二值交叉熵会让背景类主导梯度。我选用BCEWithLogitsLoss加上Dice Loss的组合class ComboLoss(nn.Module): def __init__(self, weight_bce0.5, weight_dice0.5): super().__init__() self.weight_bce weight_bce self.weight_dice weight_dice def forward(self, pred, target): bce nn.functional.binary_cross_entropy_with_logits(pred, target) pred_prob torch.sigmoid(pred) smooth 1e-6 intersection (pred_prob * target).sum() dice 1 - (2.0 * intersection smooth) / (pred_prob.sum() target.sum() smooth) return self.weight_bce * bce self.weight_dice * dice评价指标不要只看准确率因为类别不平衡时准确率没有意义。我主要看IoU交并比和F1分数。在验证集上计算IoU时需要把预测概率大于0.5的像素置为1其余置0然后和标签做对比。3.3 训练参数配置与调优训练参数如下输入尺寸512×512Batch size8如果显存有限可以降到4优化器Adam初始学习率1e-4学习率调整ReduceLROnPlateaupatience10factor0.5训练轮次80早停验证集IoU连续15轮不上升时停止我用了迁移学习将ResNet34在ImageNet上的预训练权重加载进来。遥感图像和自然图像有一定差距但底层边缘、纹理特征是通用的所以迁移学习能加速收敛。训练中前10个epoch冻结编码器只训练解码器之后解冻全部层学习率降低到5e-5进行微调。3.4 模型训练过程实录训练过程中记录loss和IoU前期loss下降很快解码器也能很快学到基本轮廓。到了中后期loss下降变缓IoU在0.72附近徘徊。这时候我做了两个调整一是加入数据增强中的色彩抖动二是把输入图像归一化。归一化的均值和标准差用ImageNet的标准值。大约到第40个epoch时验证集IoU提升到0.83最终在80轮时稳定在0.86左右。如果训练过程中出现loss反而升高的情况通常是学习率过大或者标签有错。我后来检查发现有一部分标注掩码的滑坡边界明显偏离影像中的实际滑坡体修正后IoU又涨了0.02左右。所以数据质量的重要性不亚于模型结构。4. 部署与推理实践4.1 环境配置与依赖安装项目依赖相对简单建议用conda创建独立环境conda create -n landslide python3.8 conda activate landslide pip install torch torchvision pip install opencv-python pillow numpy tqdm这里注意Pytorch版本CPU版本和GPU版本的安装命令不同。如果你有NVIDIA显卡先在终端输入nvidia-smi确认CUDA版本再去Pytorch官网选择对应的安装命令。项目里我用的Pytorch 1.12CUDA 11.6更早或更新的版本也基本兼容。如果你只有CPU环境也可以运行推理只是速度慢一些。对于单张512×512的影像CPU推理大约需要3到5秒GPU只需要0.2秒左右。所以建议有条件还是配置GPU环境。4.2 加载训练好的模型进行推理推理脚本的核心逻辑是加载权重、读图、预处理、前向传播、后处理保存结果。这里给出一个简化版import torch import cv2 import numpy as np from models.unet_resnet import ResNet34UNet model ResNet34UNet(num_classes1, pretrainedFalse) model.load_state_dict(torch.load(weights/best_model.pth, map_locationcpu)) model.eval() image cv2.imread(data/test/img_045.tif) image cv2.cvtColor(image, cv2.COLOR_BGR2RGB) image cv2.resize(image, (512, 512)) input_tensor torch.from_numpy(image).permute(2, 0, 1).float().unsqueeze(0) / 255.0 mean torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) std torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) input_tensor (input_tensor - mean) / std with torch.no_grad(): output model(input_tensor) prob output.squeeze(0).squeeze(0).numpy() mask (prob 0.5).astype(np.uint8) * 255 cv2.imwrite(result/mask_045.png, mask)如果原图尺寸大于512建议采用滑动窗口预测窗口重叠64像素最后拼接时对重叠区域取平均概率。这样得到的结果比直接resize整体预测精细很多。4.3 后处理与结果可视化模型输出的二值图会存在一些孤立小斑块和孔洞。可以做形态学开闭运算去除小于某个面积阈值的连通域。我用OpenCV的connectedComponentsWithStats统计每个连通域面积面积小于200像素的区域直接删除这样能有效减少误检。可视化时可以把预测边界叠加到原始影像上方便人工复核。我习惯生成一张叠加图作为交付物还会针对每个连通域生成一个滑坡位置列表包含中心点坐标、面积、长宽比方便后续导入GIS系统。在项目说明里我也写了如何把像素坐标换算到地理坐标前提是记录好原始影像的仿射变换参数。5. 常见问题与排查技巧5.1 显存溢出这是最常遇到的问题。512×512输入加ResNet34编码器加上batch size8需要大概8GB显存。如果你的显卡只有6GB把batch size降到4或者输入尺寸降到384×384。还有一种方法是开启梯度累积模拟更大的batch。我试过开启混合精度训练也就是Pytorch自带的torch.cuda.amp训练显存能节省30%左右速度还快了不少。不过在验证阶段还是建议用全精度避免精度损失。5.2 训练loss不下降遇到loss不下降先别急着改网络。第一步检查数据将训练样本中的影像和掩码叠加显示看看掩码是否错位、通道是否正确。第二步检查标签归一化二值掩码应该为0和1而不是0和255。第三步检查学习率学习率过大时loss会在某个值附近震荡过小时收敛缓慢。一般先用1e-4观察前10个epoch。如果数据没问题模型却一直不收敛可以尝试只用Dice Loss训练它对类别不平衡更友好而且不会像BCE那样容易陷入“全预测为背景”的局部最优。5.3 预测结果有大量椒盐噪声训练好的模型在推理时结果看起来有很多孤立像素点这通常是输入影像的拉伸方式和训练数据不匹配。遥感影像原始存储可能是16位而训练时是8位0到255归一化的。如果推理时直接读16位图像除以65535再做归一化而训练时用的是除以255输出的概率分布自然不对。解决办法是在预处理里统一数值范围。我写了一个简单的线性拉伸取影像2%到98%分位数作为最小值和最大值做归一化。这样不仅统一了尺度还增强了图像对比度预测效果会稳很多。5.4 Pytorch版本兼容性问题项目里的权重文件是用Pytorch 1.12保存的如果你用Pytorch 2.0以上版本加载有时会报UnpicklingError或者权重名称不匹配。最稳的办法是加载时加weights_onlyTrue参数新版本推荐或者用torch.load(..., map_locationcpu)避免GPU序列化问题。如果报缺少module.前缀说明训练时用了nn.DataParallel包装加载后需要去掉前缀。另一个常见问题是预训练权重下载不了。如果使用torchvision.models.resnet34(pretrainedTrue)时网络不好会导致加载卡住。我在项目里提供了下载好的预训练权重文件放在weights/目录下也可以设置环境变量TORCH_HOME指定本地缓存。6. 经验总结与后续扩展6.1 几个让我印象深刻的坑第一个坑是数据划分。最开始我按切片随机划分训练集和验证集来自同一景影像模型验证IoU高达0.91看起来很漂亮但换到另一景影像时直接掉到0.6。按区域划分后才真正反映泛化能力。第二个坑是标签栅格化时没有做地理配准导致掩码整体偏移了几个像素。这个偏差在肉眼层面几乎看不出来但会在滑坡边缘造成明显的分割错误。后来我写了一套自动检查脚本对每个切片计算掩码和影像边缘的相关系数帮助定位这类问题。第三个坑是验证集的滑坡数量太少无法稳定评估模型。后来我额外补充了一些公开的滑坡标注数据把验证集扩大到包含各种地形和植被条件评估结果才可信。6.2 可以怎么继续扩展这个项目的基础流程可以直接迁移到其他遥感分割任务比如建筑物提取、水体提取、道路提取。只需要更换数据集和调整类别数。如果想进一步提升滑坡识别精度可以从两个方向入手一是引入注意力机制比如在解码器部分加CBAM模块让网络更关注滑坡边界区域二是使用多时相遥感数据把灾前灾后影像作为双通道输入这样能识别出新发生的滑坡降低误检。另外训练好的模型可以作为预训练模型在新区数据上进行微调。遥感影像的传感器差异较大换了一个卫星数据源后可以用少量标注样本做微调通常几十张标注图就能达到还不错的精度。这个思路在项目说明里也有详细描述。最后再分享一个小技巧训练时每隔几个epoch保存一次检查点不要只保存最终模型。这样在后续调参时可以回退到特定阶段不用每次从头训练。我在项目里加了--resume参数直接指定检查点路径就能接着训练遇到停电或显存溢出也不用担心白跑。本文还有配套的精品资源点击获取
返回列表