简介面向遥感图像处理与深度学习毕业设计人群提供一套基于Python实现的高分辨率城市遥感图像水体提取完整项目方案。压缩包共二十七个文件包含十一个Python源码文件、十三张网络结构示意图、一份CSV结果数据、一篇Markdown说明文档及一个模型权重文件整体仅约七百二十六KB轻量便携部署门槛低。项目已有199人学习下载。代码注释详细适合新手阅读覆盖数据预处理、U-Net与AttU-Net模型定义、训练测试、单张预测等完整流程文档梳理目录结构并给出运行说明内置数据集可直接训练验证可支撑毕业设计、期末大作业快速落地。项目经过严格调试功能完整界面清晰可应用于城市水资源监测、洪涝灾害评估、环保规划等实际场景对学习遥感水体提取深度学习方法有较高参考价值。1. 高分辨率城市遥感图像的水体提取这份毕设源码包把深度学习落地的每一步都备齐了城市遥感影像里的水体提取是遥感方向毕业设计里最常被点名、也最容易翻车的题目。难不在网络结构而在数据——高分辨率影像里建筑物阴影、柏油路面和细小河道长得太像普通分割模型很容易把阴影当水、把水漏掉。这份以 Python 深度学习为核心的水体提取源码包把 U-Net 和 Attention U-Net 两条分割主干、数据预处理、训练评估、预训练权重 train.pth 全部打包好还带一份可直接跑的 README 和数据集适合正在做高分毕设、课程设计或者刚接触遥感图像分割但不想从零抠代码的人。它解决的核心问题是拿到一份城市遥感图怎么稳定地把水体像元分割出来并且能交得出可量化的精度指标。2. 包内结构拆解从 network.py 到 train.pth每个文件都在解决什么2.1 核心文件功能定位与调用关系拿到压缩包后先别急着跑 main.py先把目录结构捋清楚。这个包的文件命名比较直白基本按职责划分成数据、模型、训练、评估四个层次我整理了一张功能清单文件/目录作用调用方network.py定义 U-Net、Attention U-Net 网络结构solver.py、test_one_data.pydataset.pyPyTorch Dataset负责读取影像与标签data_loader.pydata_loader.py封装 DataLoader控制 batch、shufflesolver.pydata_preprocess.py影像裁剪、归一化、数据集划分手动运行一次enhance_image.py图像增强处理光照不均与阴影手动运行或接入预处理solver.py训练主循环含损失、优化器、模型保存main.pyevaluation.py计算 IoU、Dice、准确率等指标main.pytest_one_data.py单张影像预测与可视化手动运行train.pth训练好的模型权重test_one_data.py、evaluation.pyres.csv评估结果记录打开查看models/训练过程中保存的 checkpointsolver.py 写入data_set/ 与 Urben_pre/原始数据与预处理后的数据dataset.py 读取整个流程是data_preprocess.py 先把原始城市遥感图切成固定尺寸的训练块再经 enhance_image.py 做增强dataset.py 负责按路径加载影像和对应的标签掩膜solver.py 拿到 data_loader 出的批次后开始训练训练完用 evaluation.py 算指标最后 test_one_data.py 拿 train.pth 对单张图做推理。main.py 更像总控入口把上面几步串起来。2.2 U-Net 与 Attention U-Net两条分割主干的选型逻辑包内给了两张网络结构图 U-Net.png 和 AttU-Net.png对应 network.py 里的两种模型。水体提取本质是像素级二分类U-Net 是这类任务的默认起点。它由编码器、解码器和跳跃连接三部分组成编码器逐层下采样提取语义特征解码器逐层上采样恢复空间分辨率跳跃连接把编码器底层的细节纹理直接拼到解码器对应层解决下采样导致的小目标丢失问题。城市水体里的细小河流、池塘边缘靠的就是这条跳跃连接保住边界信息。Attention U-Net 在 U-Net 基础上加了注意力门控在跳跃连接前对编码器特征做加权——让网络重点看水体区域的纹理抑制建筑物阴影、柏油路面的响应。实际训练中如果用普通 U-Net 在城区影像上跑误检主要集中在阴影区域切到 Attention U-Net 后阴影误检会明显下降代价是训练时间增加 20% 左右。包内同时提供这两种结构我的建议是如果数据量少、任务以宽河面为主直接用 U-Net收敛快如果数据集中小水体多、城区阴影重用 AttU-Net 更稳。2.3 预训练权重与结果文件先看输出再决定要不要重训train.pth 是已经跑完的权重文件大约几十到上百 MB具体大小取决于网络深度。第一次拿到包不要急着删掉它重训先跑 test_one_data.py 加载这个权重对测试集里一张图做预测看输出掩膜的水体轮廓是否完整。如果效果不错说明训练流程是通的你再在此基础上做微调如果效果不理想优先检查输入影像的预处理方式和 train.pth 是否对应。res.csv 里保存的是历史评估记录一般包含每个测试样本的 IoU、Dice、准确率、召回率。打开看一眼这些数字能快速判断这个包在你这份数据上的基准水平。如果 res.csv 里的 IoU 在 0.85 以上说明原数据集上训练是充分的你后续改网络或加数据增强时以它为基准对比即可。如果 res.csv 缺失或数字偏低可能是交付者运行环境不完整建议自己重新跑一轮评估。3. 把原始遥感图喂进网络预处理、增强与数据加载的完整链路3.1 data_preprocess.py裁剪、归一化与数据集划分高分辨率城市遥感图动辄上万像素宽直接整图送进网络显存必然溢出所以第一步是裁剪。这个脚本在常见实现里做的事是按固定窗口比如 512×512把大图切块过滤掉全黑或全白的无效块再按比例划分训练集、验证集、测试集。裁剪尺寸是这里最重要的参数窗口太大会浪费计算资源窗口太小则可能切碎水体目标。# data_preprocess.py 核心逻辑常见实现 import os import numpy as np from PIL import Image from sklearn.model_selection import train_test_split crop_size 512 # 裁剪窗口根据显存调整 stride 256 # 滑动步长重叠裁剪增加样本量 image_dir data_set/images mask_dir data_set/masks image_list, mask_list [], [] for name in os.listdir(image_dir): img Image.open(os.path.join(image_dir, name)) mask Image.open(os.path.join(mask_dir, name.replace(.jpg, .png))) w, h img.size # 滑动窗口裁剪步长小于窗口尺寸制造重叠 for y in range(0, h - crop_size 1, stride): for x in range(0, w - crop_size 1, stride): crop_img img.crop((x, y, x crop_size, y crop_size)) crop_mask mask.crop((x, y, x crop_size, y crop_size)) image_list.append(np.array(crop_img)) mask_list.append(np.array(crop_mask)) # 按 7:1.5:1.5 划分训练/验证/测试 train_idx, test_idx train_test_split( range(len(image_list)), test_size0.15, random_state42 )这段代码的逻辑核心是重叠裁剪stride 取 256 而窗口是 512意味着相邻两个样本有 256 像素的重叠区域。重叠的目的是让水体边缘在多次裁剪中出现在不同位置变相增加训练样本同时避免目标被窗口边界切断。分割时固定 random_state42保证多次运行划分结果一致这点在对比实验结果时很重要——否则每次划分不同指标波动无法归因于模型改动。归一化在训练前由 dataset.py 内部完成常见做法是用 ImageNet 的均值方差做标准化。城市遥感影像虽然是 3 通道但光谱分布和自然图像差异较大如果你发现模型收敛慢可以统计自己数据集的 RGB 均值和方差替换掉 ImageNet 默认值收敛速度往往有明显提升。3.2 enhance_image.py城市阴影和水体反光的预处理策略城市水体提取的误检重灾区是阴影和水体光谱接近增强脚本主要针对这个矛盾。常见的 enhance_image.py 里包含两类操作一是对比度拉伸让水体暗部与阴影的灰度差拉大二是基于直方图均衡化的亮度校正处理城市里光照不均、高楼遮挡造成的水面亮度差异。# enhance_image.py 核心逻辑常见实现 import cv2 import numpy as np def clahe_rgb(img, clip_limit2.0, grid_size(8, 8)): 对 RGB 三通道分别做 CLAHE限制对比度避免噪声放大 clip_limit 控制对比度增强强度太大容易把噪声提亮 lab cv2.cvtColor(img, cv2.COLOR_RGB2LAB) l_channel, a_channel, b_channel cv2.split(lab) clahe cv2.createCLAHE(clipLimitclip_limit, tileGridSizegrid_size) l_clahe clahe.apply(l_channel) lab_clahe cv2.merge((l_clahe, a_channel, b_channel)) return cv2.cvtColor(lab_clahe, cv2.COLOR_LAB2RGB) def normalize_to_01(img): # 从 0-255 转到 0-1保持 float32 精度 return img.astype(np.float32) / 255.0我在跑城市水体数据时习惯把 clip_limit 设在 2.0 附近太低对阴影不敏感太高会把建筑纹理的噪声一并增强干扰分割。对水体来说真正有效的做法是对 CLAHE 后的 L 通道再叠加一次 gamma 校正gamma 取 0.8 左右可以让暗部水体更亮、与阴影的分界更明显。注意数据增强和真实推理要保持一致如果训练时加了 CLAHE测试单张图时也必须先走一遍相同的增强再进网络否则输入分布不一致指标会掉一截。3.3 Dataset 与 DataLoader训练集、验证集划分背后的细节dataset.py 是 PyTorch 的 Dataset 子类核心职责是给定一个索引返回对应的影像张量和标签张量。这里的坑在于遥感影像的标签掩膜是单通道灰度图像素值 0 表示背景、255 表示水体Dataset 里必须把这个 255 归一化成 1否则算交叉熵损失时会在 0 和 255 之间产生巨大梯度。# dataset.py 核心逻辑常见实现 from torch.utils.data import Dataset import cv2 import torch class WaterBodyDataset(Dataset): def __init__(self, image_list, mask_list, use_augmentTrue): self.image_list image_list self.mask_list mask_list self.use_augment use_augment def __len__(self): return len(self.image_list) def __getitem__(self, idx): img self.image_list[idx].astype(float32) / 255.0 mask self.mask_list[idx].astype(int64) mask[mask 255] 1 # 关键二值标签归一化 img_tensor torch.from_numpy(img).permute(2, 0, 1) mask_tensor torch.from_numpy(mask).unsqueeze(0) return img_tensor, mask_tensor标签归一化的坑我踩过不止一次如果不把 255 映射成 1模型训练初期 loss 会非常大因为 Softmax 输出的概率在 0~1 之间而标签是 255交叉熵会把两个分布拉得非常远最终导致梯度激增甚至 NaN。mask_tensor.unsqueeze(0) 是为了保持与输出张量形状一致网络输出是 [B, C, H, W]标签也必须是 [B, 1, H, W] 才能直接算损失。data_loader.py 里封装 DataLoader 时建议把 num_workers 设成 2 或 4Windows 下 num_workers 过大容易报内存错误。drop_last 参数建议打开遥感数据裁剪后样本数往往不是 batch 的整数倍最后一个 batch 不够大BatchNorm 层的统计量会抖影响训练稳定性。4. 训练与复现solver.py 里的关键参数和一次完整训练流程4.1 训练入口batch、学习率与损失函数怎么配solver.py 是训练的核心文件封装了优化器、损失函数、学习率调度和模型保存逻辑。以这个水体提取任务来看损失函数最常用的是 DiceLoss 和 BCEWithLogitsLoss 的组合因为水体类别在影像中占比往往只有 10%~20%如果单纯用交叉熵模型会趋向于把所有像素预测为背景得到很高的准确率但 IoU 很低。# solver.py 训练参数区常见实现 import torch import torch.nn as nn class Solver: def __init__(self, net, devicecuda): self.net net.to(device) self.device device # 学习率 1e-3 起步配合 StepLR 每 30 轮降 0.5 self.optimizer torch.optim.Adam(net.parameters(), lr1e-3) self.scheduler torch.optim.lr_scheduler.StepLR( self.optimizer, step_size30, gamma0.5 ) self.bce nn.BCEWithLogitsLoss() self.dice_loss self._dice_loss def _dice_loss(self, pred, target): smooth 1.0 pred torch.sigmoid(pred) intersection (pred * target).sum() return 1 - (2 * intersection smooth) / ( pred.sum() target.sum() smooth ) def train_one_step(self, image, mask): self.optimizer.zero_grad() pred self.net(image) loss self.bce(pred, mask) self.dice_loss(pred, mask) loss.backward() self.optimizer.step() return loss.item()这份代码的逻辑是BCE 负责逐像素分类的准确性DiceLoss 负责缓解类别不平衡。两者直接相加权重默认都是 1实测下来在水体提取任务上比单用其中一个稳定。学习率用 1e-3 起步配合 StepLR 每 30 个 epoch 衰减一半防止后期在最优解附近震荡。有一个参数容易被忽视BCEWithLogitsLoss 的输入是不带 sigmoid 的 logits而 dice_loss 里内部做了 sigmoid。如果训练代码里先对 pred 做 sigmoid 再传给两个损失梯度会被 sigmoid 的饱和区压缩导致训练前期几乎不更新。4.2 从 train.pth 加载权重直接预测与评估包里的 train.pth 是已经训练好的权重加载它的代码路径有两种一种是 fine-tune也就是在它基础上用小学习率继续训练另一种是直接推理也就是 evaluation.py 和 test_one_data.py 走的路。# test_one_data.py 直接预测逻辑常见实现 import torch from network import AttU_Net from PIL import Image import numpy as np device cuda if torch.cuda.is_available() else cpu model AttU_Net(img_ch3, output_ch1).to(device) # 加载权重注意 map_location 在无 GPU 机器上必须指定 state_dict torch.load(train.pth, map_locationdevice) model.load_state_dict(state_dict) model.eval() img Image.open(test_city.png).resize((512, 512)) img_tensor torch.from_numpy(np.array(img)).float().permute(2, 0, 1) / 255.0 img_tensor img_tensor.unsqueeze(0).to(device) with torch.no_grad(): pred model(img_tensor) pred torch.sigmoid(pred) pred_mask (pred 0.5).float().cpu().numpy() # 阈值 0.5 二值化这段代码输出 0/1 的水体掩膜。阈值 0.5 是默认值但如果你发现预测结果水体面积偏大或偏小可以改为 0.4 或 0.6 再做二值化这在最后的评估指标上会有明显差异。另外要注意 map_location 参数——如果你的机器只有 CPU而 train.pth 是在 GPU 上保存的不指定 map_location 会直接报显存相关错误。4.3 用 evaluation.py 和 res.csv 判断模型是否够用评估部分最常见的指标是 IoUIntersection over Union它衡量预测的水体区域与真实水体区域的交叠比例。城市水体提取的 IoU 能到 0.85 以上已经算不错低于 0.7 说明预测结果里错检漏检比较严重。# evaluation.py 指标计算核心逻辑常见实现 def compute_iou(pred, target): # pred 和 target 都是 0/1 的数组 intersection (pred target).sum() union (pred | target).sum() return intersection / (union 1e-6) def compute_dice(pred, target): intersection (pred target).sum() return (2 * intersection) / (pred.sum() target.sum() 1e-6)res.csv 里每一行通常对应一个测试样本列包括文件名、IoU、Dice、准确率、召回率。拿到 res.csv 后第一件事不是看平均 IoU而是找出 IoU 最低的几个样本看它们的影像特征。我遇到过的情况是所有低分样本都集中在窄河道或深色屋顶区域这说明模型不是整体不行而是对某类易混目标有系统性误判。这时优先针对性增强而不是盲目调网络结构。5. 常见问题与避坑训练跑飞、预测全黑、显存溢出的排查路径5.1 训练 loss 不降甚至直接变成 NaN现象solver.py 训练过程中前几个 epoch loss 输出忽大忽小随后直接打印 nan。原因最常见的是学习率过大加上标签没有归一化。如果 dataset.py 里没有把掩膜的 255 改成 1或者输入影像没有除以 255网络输出的 logits 范围会很大配合 BCE 损失会产生极大梯度。另一个常见原因是训练批次中混入了异常样本比如裁剪时没有过滤掉全为 255 的纯白图。纯白影像没有有效水体信息会让损失产生一个极大的梯度脉冲。解决先用一个 batch 的样本测试 forward 和 backward确认 loss 能降再检查 dataset.py 返回的 mask 张量取值范围是否在 0 到 1 之间最后把学习率降到 1e-4 试跑。如果 loss 恢复正常再逐步调回去。5.2 预测结果全黑或全白现象test_one_data.py 输出的掩膜一片 0 或一片 1看不到水体轮廓。原因出现全黑通常是模型把输入预测成全部背景常见于训练时类别不平衡未处理或者阈值得到了 1.0 以上出现全白常见于网络输出 logits 全部为正值且阈值设得太低。还有一个隐蔽原因是输入影像的通道顺序不对——OpenCV 读进来是 BGR而训练时 PIL 读的是 RGB通道错位会让模型看到完全不同的图像。解决把预测前的 sigmoid 输出统计一下 min、max、mean。如果 mean 接近 0.5 但二值化后全黑说明阈值设太高改成 0.3 再试如果 sigmoid 前的 logits 全部是正数说明模型偏向把所有东西判为水体这时需要检查训练损失里是否只用了 DiceLoss因为 Dice 对背景区域的惩罚不够通常要和 BCE 相加。5.3 GPU 显存溢出现象训练或预测时报 CUDA out of memory。原因遥感影像裁剪尺寸太大或 batch size 偏大。很多同学直接把 1024×1024 的块送到 Attention U-Net 里跑注意力模块额外增加的中继特征在显存里堆叠特别容易爆。解决把 crop_size 从 512 缩到 384 或 256batch size 从 8 降到 4。如果必须用大图预测走第六章的滑动窗口方案。另外solver.py 里的梯度累积技巧也可以解决这个问题不减少 batch而是每 4 个 step 再更新一次梯度等效于 batch size 扩大 4 倍显存只占原来的四分之一。5.4 水体提取结果与建筑物阴影混淆现象模型的预测掩膜里大片阴影区域被标成水体和真实水体连成一片。原因高分辨率遥感图里建筑物阴影的光谱特征与深色水体非常接近RGB 三通道下人眼都很难区分。只靠 RGB 图像做二分类上限就在这里。解决优先检查 enhance_image.py 里 CLAHE 的参数是否生效clip_limit 调低会加剧混淆。如果增强已做但还是分不清后续可以引入近红外波段或者把 Attention U-Net 的注意力权重可视化确认模型是否真的聚焦在水体区域。这个包是纯 RGB 输入阴影混淆是它的固有边界遇到难分样本时可以通过后处理比如叠加建筑高度数据做掩膜来修正。5.5 加载 train.pth 报键名不匹配现象main.py 加载权重时报 missing keys 或 unexpected keys。原因network.py 里选择的网络结构和训练时不一致。比如权重文件是用 Attention U-Net 训练的你加载到 U-Net 上键名自然对不上。解决确认加载权重的网络类和训练时完全一致包括 img_ch 和 output_ch 参数。如果包内有两个网络定义加载前先打印 state_dict 里的键名和当前网络的键名对比不一致就换网络结构。6. 进阶技巧用滑动窗口预测大幅影像并输出矢量化边界当 test_one_data.py 只能处理 512×512 的输入时面对整幅城市遥感大图会失效。我的做法是写一个滑动窗口推理函数按 stride 在大图上切块预测再把结果拼接回完整的坐标空间。拼接时每个窗口的预测置信度可以叠加重叠区域取平均这样边界过渡比硬裁剪自然得多。# 滑动窗口推理核心逻辑 def predict_full_image(model, full_img, crop_size512, stride256, devicecuda): h, w full_img.shape[:2] pred_map np.zeros((h, w), dtypenp.float32) count_map np.zeros((h, w), dtypenp.float32) for y in range(0, h - crop_size 1, stride): for x in range(0, w - crop_size 1, stride): crop full_img[y:y crop_size, x:x crop_size] crop_tensor torch.from_numpy(crop).permute(2, 0, 1).unsqueeze(0) with torch.no_grad(): out torch.sigmoid(model(crop_tensor.to(device))).cpu().numpy() pred_map[y:y crop_size, x:x crop_size] out[0, 0] count_map[y:y crop_size, x:x crop_size] 1 pred_map pred_map / np.maximum(count_map, 1) return pred_map拼接后用形态学后处理收尾先做一次开运算剔除孤立小斑块再做一次闭运算填补河流中间的空洞。如果到这一步还需要交矢量成果可以对二值掩膜调用 OpenCV 的 findContours 提取轮廓导出成 Shapefile 或 GeoJSON。边界平滑这块建议在矢量化后做对掩膜做高斯模糊再二值化容易让细碎锯齿变少。这类水体提取项目在答辩环节最常被追问的问题是“你的方法在阴影区域为什么没有漏检”。我的习惯是保存每一轮训练的 checkpoint在镜像里挑一张阴影密集的图把注意力可视化图、预测掩膜、原图三张图拼在一起放进论文附录效果比写一大段文字好得多。从那以后我每跑一个分割实验都会强制走一遍“训练保存中间权重 难例可视化 量化指标”这三步闭环。这个包的代码路径已经很顺顺着跑通再在数据增强和后处理上做点自己的实验希望帮到你。本文还有配套的精品资源点击获取
