简介:城市遥感图像水体识别实战包面向需要开展遥感分割实验的本科生与开发者,基于高分辨率影像提供U-Net与AttU-Net两种主流结构的完整实现。资源共32个文件,以11个Python脚本、13张示例图像、3个CSV标签文件为主,整体压缩包仅733KB,并附带已训练权重train.pth与数据增强脚本,可在Python 3.8+、PyTorch 1.9+环境直接运行。数据集Urban_pre及res.csv支持快速验证,训练、评估、单图测试模块划分清晰,从数据预处理到结果导出全流程覆盖。已有133人学习下载,适合作为本科毕业设计、遥感课程大作业或轻量级水体监测原型的起步框架。 最近在做一个城市水体提取项目,把一套相对完整的实战代码整理了出来:城市遥感图水体识别实战包,包含PyTorch模型、标注数据集和完整的训练推理代码。简单说,就是给你一批城市遥感影像,通过深度学习模型自动把河流、湖泊、池塘这些水体区域分割成二值掩膜。无论你是做遥感、GIS、还是刚入门CV分割,这套东西都能直接跑通,省去从零搭环境、标数据、调模型的重复劳动。
水体识别这事儿看着简单,做起来坑不少。城市环境下阴影、高楼、植被、柏油路都会干扰,裸用传统阈值法根本压不住。用PyTorch做语义分割是目前比较稳的路线,关键是数据集和训练推理的工程细节。这篇文章不聊虚的,就把我当时从数据准备、模型选型、训练调参到推理后处理的完整思路和踩坑记录写出来,代码和配置都以这套实战包为准。
1. 项目整体思路与选型考量
1.1 为什么用PyTorch而不是其他框架
水体识别本质上是语义分割任务,输入一张遥感影像,输出每个像素是水体还是背景。PyTorch在这个领域的使用率确实高,主要原因有三点:一是生态里现成的分割模型多,比如Unet、DeepLabV3+、SegFormer,PyTorch的实现都比其他框架更全;二是调试直观,动态图机制遇到维度不对或者loss异常,直接print就能定位;三是社区活跃,遥感方向的预训练权重、数据增强方案、推理trick大部分都优先发PyTorch版本。
我之前也试过TensorFlow,但做自定义Dataset和数据增强时总觉得绕。PyTorch的Dataset和DataLoader接口写起来顺手,尤其是做遥感大图推理时,需要用滑动窗口切图,配合torchvision.transforms做在线增强,逻辑很清晰。如果你的机器是NVIDIA显卡,PyTorch配合CUDA的安装也很成熟,实测下来比TF省心。
1.2 数据集的构建思路和标注难点
这套实战包里的标注数据集是我从公开城市遥感影像中挑选并手工标注的,覆盖了城市河流、公园湖泊、人工水渠、蓄水池这几类典型水体。标注格式是常见的语义分割掩膜,背景为0,水体为1,使用LabelMe画多边形后转成png灰度图。
城市水体标注有几个容易翻车的点。第一个是阴影:高楼投下的阴影和水体在光谱上很像,肉眼都容易看错,标注时我定了一条规则——拿不准的阴影区域宁可标背景,也不要模棱两可地标成水体,否则模型会被带偏。第二个是细小水体:城市里的窄水渠、小喷泉在低分辨率影像上很可能只有几个像素,若强行标注会让训练标签很碎,模型学不到有效特征,所以我只保留宽度大于5个像素的水体。第三个是水体边界:部分水体边缘有植被遮挡或者码头建筑,边界线条不好画,我统一用影像上肉眼可见的水陆分界线为准,保证标注一致性。
1.3 模型选型:从Unet到轻量化变体
模型这块我一开始直接上了Unet,因为结构简单、适合中小规模数据集。Unet的编码器-解码器结构能有效融合多尺度特征,对水体这种边缘相对平滑的目标来说足够用。但Unet原版在训练时占用显存偏高,于是我换成了以ResNet34为backbone的Unet变体,分割头不变,编码器换成预训练模型,速度和精度平衡了不少。
后来我还尝试了DeepLabV3+,ASPP模块能进一步扩大感受野,对大面积水体效果更好,但训练轮数需要更多。如果你追求更高精度,可以考虑SegFormer这类Transformer结构,但对城市遥感这种场景来说,模型复杂度上去了,推理速度下来了,实测收益不大。最终包里的主模型还是ResNet34+Unet,兼顾精度和推理效率,显存占用也友好。
2. 数据准备与预处理细节
2.1 标注格式转换与数据集目录结构
用LabelMe标注完会生成JSON文件,里面是多边形坐标,不能直接喂给模型。我写了一个转换脚本,把JSON里的多边形转成掩膜png,核心是用PIL的ImageDraw.Polygon填充。这里有一个关键步骤:原始遥感影像可能是RGB三通道,标注掩膜一定是单通道灰度图,不要存成RGB三通道的png,否则后续dataloader会多出两个维度,训练直接报错。
数据集的目录结构建议固定下来,方便复用:
dataset/ ├── images/ │ ├── img_001.tif │ ├── img_002.tif │ └── ... ├── masks/ │ ├── img_001.png │ ├── img_002.png │ └── ... └── train_val_split.txtimages里放原始影像,masks里放对应的掩膜,两者文件名一一对应。train_val_split.txt每一行是“图片名 标签”,我会按8:2的比例划分训练集和验证集,划分时注意把同一片区域相近的影像尽量放到同一边,避免验证集和训练集场景太像,导致验证指标虚高。
2.2 影像裁剪与数据增强策略
遥感影像分辨率很高,一般不能整图直接进模型。我把影像裁剪成256x256的小块,裁剪方式采用重叠裁剪:步长设为192,这样相邻图块之间有64像素的重叠,能有效减少推理时拼接的边缘不连续问题。训练和验证统一裁剪,推理时也用同样的滑动窗口方式处理。
数据增强我用了随机水平翻转、垂直翻转、旋转90度、随机亮度对比度调整。这里要注意,水体识别对颜色比较敏感,亮度调整幅度不要太大,我在实战里把brightness和contrast范围控制在0.8到1.2之间,太强的颜色抖动会把水体本身的深色特征破坏掉,反而降低精度。另外,遥感影像里植被和水体在近红外波段差异很大,如果只有RGB影像,就不要试图通过颜色增强来模拟多光谱信息,模型会学到错误的颜色关联。
2.3 Dataset类的实现细节
PyTorch的Dataset类需要自己写,代码不复杂,但有几个细节值得注意。第一,读取影像用rasterio或cv2都行,但要注意通道顺序,cv2读取出来是BGR,需要转成RGB。第二,掩膜读取后要确保是0和1两种值,有些标注工具导出的掩膜可能是0和255,需要除以255归一化。第三,数据增强需要同时作用于影像和掩膜,不能只对影像做增强而掩膜不变,否则模型学不到对应关系。
我贴一下核心代码思路:
class WaterDataset(Dataset): def __init__(self, image_dir, mask_dir, file_list, size=256, transform=None): self.image_dir = image_dir self.mask_dir = mask_dir self.file_list = file_list self.size = size self.transform = transform def __len__(self): return len(self.file_list) def __getitem__(self, idx): name = self.file_list[idx].strip() img = cv2.imread(os.path.join(self.image_dir, name)) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) mask = cv2.imread(os.path.join(self.mask_dir, name.replace('.tif', '.png')), cv2.IMREAD_GRAYSCALE) mask = (mask > 127).astype(np.uint8) # 这里统一resize到self.size,注意mask用最近邻插值 img = cv2.resize(img, (self.size, self.size), interpolation=cv2.INTER_LINEAR) mask = cv2.resize(mask, (self.size, self.size), interpolation=cv2.INTER_NEAREST) # ... return torch.from_numpy(img).permute(2,0,1).float() / 255.0, torch.from_numpy(mask).long()掩膜resize一定要用最近邻插值,不要用线性插值,否则掩膜边缘会出现介于0和1之间的小数,导致训练时交叉熵损失出现莫名其妙的异常值。
3. 模型训练全流程实操
3.1 环境搭建:PyTorch+CUDA配置
环境这块我给一个可以直接照着装的版本组合:Python 3.10 + PyTorch 2.0.1 + CUDA 11.8,配上torchvision 0.15.1。这几项实测下来兼容性很好,能覆盖大多数主流显卡。
安装命令我习惯用conda先建虚拟环境:
conda create -n water python=3.10 conda activate water pip install torch==2.0.1 torchvision==0.15.1 --index-url https://download.pytorch.org/whl/cu118如果你是NVIDIA 30系或40系显卡,CUDA 11.8这个组合基本没问题;如果是A卡,PyTorch的ROCm支持也可以,但配置起来麻烦一些,建议直接用CPU先跑通流程,再考虑加速。安装完成后一定要验证CUDA是否可用:
import torch print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))如果输出False,多半是PyTorch版本和显卡驱动不匹配,或者装了CPU版本的torch。我之前因为下载速度慢,图省事装了CPU版,跑训练才发现速度慢到无法忍受,后来还是老老实实重装了CUDA版。
3.2 训练脚本核心代码解析
训练脚本我分成几个模块:数据加载、模型定义、损失函数、优化器、训练循环、验证循环。整体流程其实很固定,关键是几个容易被忽略的细节。
模型部分使用segmentation_models_pytorch库,封装得很好:
import segmentation_models_pytorch as smp model = smp.Unet( encoder_name="resnet34", encoder_weights="imagenet", in_channels=3, classes=1, )这里我把classes设为1,配合带logits的输出,训练时用BCEWithLogitsLoss。如果classes设为2,就要配合CrossEntropyLoss,输出通道变成2,效果差异不大,但二分类用单通道logits更节约显存。
优化器用AdamW,初始学习率1e-4,weight_decay设为1e-4。训练循环里每一个epoch结束后在验证集上计算IoU,如果IoU连续10个epoch不涨就提前终止。我贴一下损失函数组合的核心逻辑:
bce = nn.BCEWithLogitsLoss() dice = smp.losses.DiceLoss(mode='binary') bce_weight = 0.5 dice_weight = 0.5 for images, masks in train_loader: images = images.to(device) masks = masks.to(device).float() logits = model(images) loss = bce_weight * bce(logits, masks.unsqueeze(1)) + dice_weight * dice(logits, masks.unsqueeze(1)) optimizer.zero_grad() loss.backward() optimizer.step()3.3 损失函数与评价指标的选择
水体识别是典型的类别不平衡问题:城市影像里水体像素占比往往不到10%,如果只用交叉熵损失,模型很容易把所有像素预测成背景,因为这样损失已经很低了。解决办法是组合损失,我最常用的是BCE Loss加Dice Loss,Dice Loss对前景占比小的目标更敏感,能有效提升水体区域的重合度。
评价指标我没有只看准确率,而是以IoU和F1为主。准确率在水体像素占比低时没有参考意义,即使全部预测为背景,准确率也可能超过90%,但这样的模型实际不可用。我在训练时会打印每个epoch的mIoU、F1、Precision、Recall,重点关注mIoU和F1,对水体边界的小目标,F1比IoU更能反映实际效果。
3.4 训练调参与检查点保存
训练过程的调参经验直接说结论。学习率方面,我用了warmup加CosineAnnealing:前5个epoch从1e-5逐步升到1e-4,之后按余弦曲线衰减到1e-6,相比固定学习率收敛更稳,最终IoU提升大概2到3个百分点。batch size在8到16之间比较合适,太大容易显存溢出,太小则BN层统计不稳定。
模型检查点保存我一般每个epoch都会存一份,但只保留验证集IoU最高的那份,避免磁盘爆炸。这里提一下“模型检查器”的概念:我写了一个简单的回调函数,每轮验证后对比最优IoU,如果连续5个epoch不提升,就把学习率乘以0.1,如果连续10个epoch不提升就直接停止训练,省得挂机跑一堆无效轮次。
训练完成后,保存整个模型或者只保存state_dict都行,我建议保存state_dict加配置信息,因为后续推理要载入时能知道当时的输入尺寸、类别数等参数:
torch.save({ 'model_state_dict': model.state_dict(), 'encoder_name': 'resnet34', 'classes': 1, 'input_size': 256, }, 'best_model.pth')4. 推理部署与效果测试
4.1 滑动窗口推理与图像后处理
训练好的模型只能处理256x256的小图,但实际遥感影像往往是几千乘几千的大图,直接resize会造成严重的信息丢失。我的做法是用滑动窗口把大图切成小块,逐个推理后拼接回原图尺寸。这里“滑动窗口滤波模型”的思路很关键:相邻窗口之间要有重叠,比如步长设为224,窗口尺寸256,这样每个窗口有32像素的重叠。推理时对重叠部分取平均值,能消除拼接缝处的像素突变。
拼接之后的二值掩膜还需要做后处理,否则会有很多离散的小噪点和空洞。我用了两步:先用中值滤波,核大小设为5,能去除孤立噪点;再用形态学开运算去除小连通域,闭运算填充小孔。后处理对最终视觉效果提升很大,有时候mIoU只涨了1个点,但图上明显干净很多。
后处理代码示例:
import cv2 import numpy as np def post_process(mask, kernel_size=5): mask = cv2.medianBlur(mask, kernel_size) kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3)) mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=1) mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=1) return mask4.2 模型导出与批量预测
推理部署时我不建议直接用训练时的PyTorch模型做在线预测,因为每次都要加载torch和模型定义,麻烦而且启动慢。我习惯先把模型导出成TorchScript或者ONNX,前者在PyTorch环境中直接加载,后者可以对接OpenCV、ONNX Runtime甚至C++部署。
导出TorchScript的代码很简单:
model.eval() dummy_input = torch.randn(1, 3, 256, 256) traced_model = torch.jit.trace(model.cpu(), dummy_input) traced_model.save('water_model.pt')批量预测时直接加载water_model.pt,读入整张大图,用滑动窗口切块加后处理,输出水体掩膜。我在实战包里还写了一个批处理脚本,支持输入一个影像文件夹,输出每个影像对应的水体掩膜和叠加效果图,方便快速查看。
4.3 实际效果评估与误差分析
我在验证集上的结果大致是:mIoU 0.82,F1 0.89,整体效果可用,但不是所有场景都理想。大水面(河流主河道、大湖)分割很干净,边缘也比较平滑;小水体(水渠、小水池)偶尔会漏检,主要因为裁剪后目标太小,特征不明显。阴影遮挡区域误检率偏高,尤其在高楼密集的城区,深色阴影和水体在视觉上确实难以区分。
针对误差,我尝试过给模型增加一个NDWI先验通道,也就是用绿波段和近红外波段计算归一化差异水体指数。如果你的遥感影像是多光谱数据,把NDWI作为第四个通道输入,误检率会明显下降;但RGB影像没法这么操作。另一个策略是做模型融合,把Unet和DeepLabV3+的输出取平均,能提升整体稳定性,但训练和推理时间增加,我最后没有在正式版里采用。
5. 踩坑记录与常见问题排查
5.1 显存溢出和数据加载瓶颈
跑训练最常见的坑是OOM。我一开始把batch size设为32,输入256x256,ResNet34编码器直接爆显存。解决办法不是硬调小batch size,而是先降低输入尺寸到224,再把batch size降到16,配合混合精度训练,显存占用能降一半。PyTorch自带的torch.cuda.amp很好用,训练速度还能提升一截。
数据加载慢也容易被忽略。我一开始直接从机械硬盘读图,训练每个epoch要等很久。后来把数据转成内存映射格式或者用SSD,并设置DataLoader的num_workers为4,prefetch_factor调成2,训练速度明显提升。注意Windows下num_workers大于0有时会报错,需要在主程序里加上ifname== 'main'保护。
5.2 水体边界模糊和阴影误检处理
水体边界是最头疼的问题。遥感影像中水体边界常常被树木、芦苇、码头遮挡,模型预测出来的边界会呈锯齿状,有时候还会把岸边湿润的泥地识别成水。我处理这类问题有三个手段:一是后处理时增加形态学闭运算的迭代次数,让边界更平滑;二是在损失函数中加入Boundary Loss或者把Dice Loss的权重提高,强化边界区域的学习;三是训练时把标注掩膜往外稍微膨胀1到2个像素,给边界一点容错空间。
阴影误检比较难处理,单纯调参效果有限。我最后的做法是加入一条人工后处理规则:对预测为水体的区域,统计其亮度均值,如果过低(接近纯黑),就判定为阴影并去除。这个方法看起来“不深度学习”,但实测能把阴影误检率降低不少,尤其在建筑物密集的区域。
5.3 不同分辨率遥感图的适配问题
训练集影像和实际预测影像分辨率不一致时,模型效果会明显下降。比如训练时用的是0.3米分辨率,推理时换成了1米分辨率的影像,物体尺度完全不同,水体边缘特征也变了。解决办法是推理前把输入影像重采样到和训练集接近的地面分辨率,而不是简单resize到固定尺寸。
如果你不确定原本分辨率,建议在推理脚本里记录影像的元数据(GSD),然后统一重采样到0.5米左右。重采样的插值方法用双线性即可,不需要太高级,因为模型本身对分辨率有一定兼容性,但差距太大会明显掉点。
5.4 常见报错与解决方案速查
我在开发过程中记录了一些高频报错,按场景整理成表格,可以帮你快速定位:
| 报错现象 | 可能原因 | 解决方案 |
|---|---|---|
| RuntimeError: CUDA out of memory | batch size太大或输入尺寸太大 | 降低batch size、输入尺寸,开启混合精度 |
| ValueError: Target and input size mismatch | 损失函数输入输出维度不匹配 | 检查mask是否少了通道维度,使用unsqueeze |
| FileNotFoundError: Can't open image file | 数据集路径或文件名不匹配 | 检查names.txt里的文件名和实际路径是否一致 |
| Loss不下降 | 学习率太高或数据增强过强 | 降低学习率,关闭颜色增强测试 |
| 验证集IoU很高但测试效果差 | 训练集和验证集分布有重叠 | 重新划分数据集,按区域划分而非随机划分 |
| DataLoader worker进程意外退出 | Windows下num_workers问题 | 在main函数入口调用multiprocessing.freeze_support |
这套实战包我后来又接到了多个城市的影像上测试,发现一个共性问题:每个城市的水体形态、背景地物都不一样,直接用预训练权重零样本推理肯定不行。最稳的做法是针对目标城市收集少量影像做微调,哪怕只有几十张标注样本,效果都能提升不少。这也是我把数据集、模型、训练推理代码打包在一起的原因——拿到手先跑通,再基于自己的数据微调,就不用来回折腾基础设施了。
如果要从头自己标注一套城市水体数据,我的建议是先标一个区域,跑通全流程,再逐步扩充。模型选Unet起步完全够用,别一开始就上多复杂的结构。训练时多盯验证集IoU和可视化结果,不要只盯着训练损失。推理阶段的后处理别省,几行形态学操作能让输出质量上一个台阶。这套包里的代码都是按这个思路写的,你照着跑一遍,对水体识别整个技术链路就心里有数了。
本文还有配套的精品资源,点击获取