深度学习遥感图像水体提取：基于U-Net与Attention U-Net的语义分割实践

简介：面向高分辨率城市遥感图像的水体提取任务，这是一套基于Python深度学习的完整毕设项目，适合作为毕业设计、期末大作业或课程设计参考。项目代码注释详细，新手也能理解，部署简单即可运行。资源包共27个文件，大小仅726KB，包含11个Python脚本（主程序、网络模型、数据加载、预处理、测试等）、13张示意图、1个预训练权重文件、1个说明文档和1个CSV结果文件，目录结构清晰，便于按模块检索学习。目前已有199人学习下载。代码涵盖U-Net/AttU-Net模型实现、数据预处理、训练与评估流程，并附有训练好的权重，可直接用于水体提取实验或在此基础上二次开发。项目经严格调试，功能完整，系统界面友好，对完成高分毕设具有较高的参考价值。

1. 高分影像水体提取，最难的不是“水”而是“像水的东西”

城市高分辨率遥感影像里的水体提取，表面看是“把蓝色像素挑出来”，真正做下去才发现误检源全在阴影、黑色屋顶、立交桥下暗部和湿润地面。这份基于 Python 深度学习的高分辨率城市遥感图像水体提取项目，把 U-Net 和 Attention U-Net 两套语义分割网络、数据预处理、增强、训练、评估和单图推理全部串成了一条完整流水线，代码带详细注释，还附带数据集与训练好的权重文件，属于直接下载就能复现的毕设级工程。适合正在做遥感语义分割毕设、课程设计，或者想快速拿一套基线代码做对比实验的从业者——你要的其实不是“理论多深”，而是一条能跑通、能改、能写进论文的完整链路，这份资源正好落在这个位置。

2. 数据准备与预处理：先把遥感影像切成模型能吃的样子

遥感影像和自然图像最大的差别在于尺寸和通道含义。一副高分影像动辄几千乘几千像素，直接送进 U-Net 会直接撑爆显存；而且影像的像素值范围、标注格式都跟日常用到的 VOC 数据集不一样，不做预处理根本跑不起来。这一章围绕项目的data_set、util/data_preprocess.py和enhance_image.py三个部分展开。

2.1 项目文件清单与目录结构

解压后的目录结构基本就是一条完整的深度学习工程线，我拆开看了一遍，核心文件对应关系如下：

2.2 预处理脚本：裁剪与归一化的两个关键设定

data_preprocess.py的核心逻辑是把大尺寸遥感影像按滑动窗口切成小 patch。常见做法是设 patch 为 256×256 或 512×512，步长与 patch 相同（不重叠），这样既保留局部空间特征，又能在单张 11GB 显存显卡上跑得动。先看一眼典型实现：

import numpy as np import cv2 PATCH_SIZE = 256 STRIDE = 256 def sliding_crop(image, label, patch_size=PATCH_SIZE, stride=STRIDE): crops = [] label_crops = [] h, w = image.shape[:2] 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] label_patch = label[y:y+patch_size, x:x+patch_size] crops.append(img_patch) label_crops.append(label_patch) return crops, label_crops def normalize_image(patch): patch = patch.astype(np.float32) / 255.0 mean = np.array([0.485, 0.456, 0.406]) std = np.array([0.229, 0.224, 0.225]) for i in range(3): patch[:, :, i] = (patch[:, :, i] - mean[i]) / std[i] return patch

滑动裁剪的逻辑不复杂：从影像左上角开始，按stride向右向下移动，每次取patch_size × patch_size的区域，影像和 label 用相同坐标裁剪，保证像素级对齐。normalize_image里用的是 ImageNet 的均值和标准差，遥感影像虽然是多波段，但 RGB 三通道用这组参数做标准化在工程实践里很常见，迁移学习场景下尤其合适。需要注意：步长设成和 patch 相等会丢掉边缘区域，如果影像尺寸不是 patch 的整数倍，会在右下角出现裁不满的情况。我一般会先对影像做边缘填充（padding），或者把步长改成 patch 的一半做重叠裁剪，后者还能顺带做数据增广，代价是样本量变大、训练变慢。

2.3 标签处理：mask 的像素值必须在训练前查一遍

水体提取本质是二分类语义分割，标签 mask 里的像素值必须严格约定：水体区域为 1，非水体为 0，整张 mask 是单通道灰度图。这个环节最容易翻车，很多项目在数据准备阶段没检查标签，训练时 loss 一直震荡，最后发现 mask 里混入了 255、128 之类的中间值。建议在预处理后立即做一次唯一值检查：

def check_label_unique(label_path): import cv2 mask = cv2.imread(label_path, cv2.IMREAD_GRAYSCALE) values = np.unique(mask) print(f"unique values: {values}") assert set(values.tolist()).issubset({0, 1}), \ f"mask values should be 0/1, got {values.tolist()}"

这段代码强制要求 mask 只含 0 和 1。如果发现 255，通常是标注软件导出时把前景标成了 255，需要做一次二值化：mask = cv2.threshold(mask, 127, 1, cv2.THRESH_BINARY)[1]。数据集的Urben_pre目录里放的应该就是预处理后的影像和 mask，直接用就行，但自己新加数据时这个检查一定要留在预处理脚本里，不要想着后面再补。

2.4 增强脚本：几何变换同步，色彩变换独立

enhance_image.py的作用是在训练时对 patch 做在线增强。核心原则是：几何变换（翻转、旋转、缩放）必须对影像和 mask 同步做同样的变换，否则标签和图像会错位；色彩变换（亮度、对比度、饱和度）只对影像做，mask 不参与。错误示范是把整张图旋转后忘了转 mask，训练时 loss 能降，但 mask 全乱，模型的边界预测会完全失真。

import random import cv2 import numpy as np def enhance_pair(img, mask): # 几何变换：同时作用于图像和 mask if random.random() > 0.5: img = cv2.flip(img, 1) mask = cv2.flip(mask, 1) if random.random() > 0.5: img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE) mask = cv2.rotate(mask, cv2.ROTATE_90_CLOCKWISE) # 色彩变换：只作用于图像 if random.random() > 0.5: hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) hsv[:, :, 2] = np.clip(hsv[:, :, 2] * random.uniform(0.8, 1.2), 0, 255) img = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR) return img, mask

这段增强策略对城市水体场景是有针对性的。翻转和旋转能提升模型对水体走向的旋转不变性；亮度扰动模拟不同光照条件下水面反射差异。要注意遥感影像的旋转角度不要随意设大角度，90 度整数倍旋转对语义分割最安全，45 度任意角度会引入插值噪声，水体边缘的锯齿反而干扰训练。

3. 模型选型与网络设计：U-Net 还是 Attention U-Net

网络结构是这套项目的核心。network.py里同时实现了 U-Net 和 Attention U-Net，前者是基线，后者是带注意力门的改进版本。两个模型共用数据集和训练流程，切换训练时只需改main.py里的模型名参数，这个设计对毕设实验对比非常友好。

3.1 为什么遥感水体分割默认选 U-Net

U-Net 的结构是编码器-解码器加跳连接。编码器逐层下采样，空间分辨率从 256 降到 32，每层提取的特征从边缘纹理逐步抽象到语义类别；解码器逐层上采样恢复分辨率；跳连接把编码器同层的高分辨率特征直接拼到解码器上。这个设计正好解决了水体提取的两个痛点：水体边界需要高分辨率特征来精确定位；水体内部需要深层语义来区分“水面”和“阴影中的深色地面”。相比 DeepLab 系列的空洞卷积，U-Net 在中小尺寸数据集上更容易训练收敛，参数量也更小，CPU 都能跑推理。

3.2 Attention U-Net 的注意力门：到底加了什么

Attention U-Net 是在解码器每个 stage 的上采样前加了一个注意力门（Attention Gate，AG）。普通 U-Net 的跳连接是把编码器特征直接拼接，不管这个位置是否有用；注意力门会先计算一个注意力权重，让解码器过滤掉无关背景区域的浅层特征，保留与水体相关的细节。在network.py里典型的注意力门实现长这样：

import torch import torch.nn as nn class AttentionGate(nn.Module): def __init__(self, in_channels, out_channels): super(AttentionGate, self).__init__() self.W_g = nn.Conv2d(in_channels, out_channels, kernel_size=1) self.W_x = nn.Conv2d(in_channels, out_channels, kernel_size=1) self.psi = nn.Conv2d(out_channels, 1, kernel_size=1) self.relu = nn.ReLU() self.sigmoid = nn.Sigmoid() def forward(self, g, x): # g: 解码器的门控信号（高层语义特征） # x: 编码器的跳连接特征（低层空间细节） theta_g = self.W_g(g) theta_x = self.W_x(x) f = self.relu(theta_g + theta_x) attn = self.sigmoid(self.psi(f)) return x * attn

这里的g来自解码器上采样后的高层特征，x来自编码器同层跳连接。注意力权重由两者的融合特征经过 1×1 卷积和 sigmoid 得到，值域在 0 到 1 之间，乘到x上就是告诉网络：这个位置如果与水体语义相关，保留细节；如果是误检高发区（比如深色屋顶），权重压低。在实际测试中，Attention U-Net 在阴影误检和细长水体连通性上通常比 U-Net 高 2 到 5 个百分点的 IoU，代价是训练速度稍慢，loss 收敛需要更多 epoch。

3.3 损失函数的选择与正负样本不均衡

水体提取里“水”的面积占比往往只有 5% 到 20%，直接拿普通交叉熵训练，模型会倾向于把所有像素预测为背景。项目里常见的做法是 Dice Loss 和交叉熵的组合加权。Dice Loss 直接优化目标区域的重叠度，适合小目标分割；交叉熵提供稳定的梯度信号，避免 Dice Loss 在极端情况下梯度异常。组合公式通常是：

import torch.nn as nn def dice_loss(pred, target, smooth=1.0): pred = torch.sigmoid(pred) intersection = (pred * target).sum() return 1 - (2.0 * intersection + smooth) / (pred.sum() + target.sum() + smooth) def combined_loss(pred, target, bce_weight=0.5, dice_weight=0.5): bce = nn.BCEWithLogitsLoss()(pred, target) dice = dice_loss(pred, target) return bce_weight * bce + dice_weight * dice

组合损失的bce_weight和dice_weight在项目里一般各设 0.5。如果发现训练初期 loss 下降缓慢，可以适当调高bce_weight；如果发现后期边缘不精细，可以调高dice_weight到 0.7。注意 Dice Loss 在预测概率接近 0 或 1 时梯度会变小，不适合单独使用，务必和 BCE 组合。

4. 训练流水线：从 dataset.py 到 solver.py 的完整闭环

这一章把dataset.py、data_loader.py、solver.py和main.py串起来。这套工程的训练流程写得很规整，数据加载、训练器、入口脚本各司其职，换数据集时只需要改dataset.py里的路径和 mask 读取方式。

4.1 dataset.py：数据加载器的关键实现

dataset.py继承 PyTorch 的Dataset类，关键点是__getitem__返回的是已经做增强、张量化、归一化的(img_tensor, mask_tensor)对。工程里典型的实现方式：

import torch from torch.utils.data import Dataset import cv2 import numpy as np from enhance_image import enhance_pair class WaterBodyDataset(Dataset): def __init__(self, image_dir, mask_dir, transform=True): self.image_paths = sorted(image_dir.glob("*.png")) self.mask_paths = sorted(mask_dir.glob("*.png")) self.transform = transform def __len__(self): return len(self.image_paths) def __getitem__(self, idx): img = cv2.imread(str(self.image_paths[idx])) mask = cv2.imread(str(self.mask_paths[idx]), cv2.IMREAD_GRAYSCALE) if self.transform: img, mask = enhance_pair(img, mask) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = img.astype(np.float32) / 255.0 img = torch.from_numpy(img).permute(2, 0, 1) mask = mask.astype(np.float32) mask = torch.from_numpy(mask).unsqueeze(0) return img, mask

这里所有变换都在 CPU 上完成，BGR2RGB是因为 OpenCV 读图默认是 BGR，而预训练模型是按 RGB 训练的。permute(2, 0, 1)把 HWC 转成 CHW，unsqueeze(0)给 mask 增加通道维度，大小变成(1, H, W)，与模型输出的(1, H, W)保持一致，才能直接算 BCELoss。这里有个容易忽略的细节：数据增强用的是enhance_pair，在线增强的好处是每个 epoch 看到的样本都不同，等于隐式扩大了数据集规模，对遥感小数据集来说作用比离线增强更明显。

4.2 solver.py：训练器封装与模型保存逻辑

solver.py是整个训练过程的中枢。它接收网络、数据加载器、优化器和超参数，在内部循环里完成前向、反向、验证和模型保存。核心训练循环大致如下：

import torch import torch.optim as optim class Solver: def __init__(self, model, train_loader, val_loader, lr=1e-4): self.model = model self.train_loader = train_loader self.val_loader = val_loader self.optimizer = optim.Adam(model.parameters(), lr=lr) self.scheduler = optim.lr_scheduler.ReduceLROnPlateau( self.optimizer, mode="min", factor=0.5, patience=8) def train_one_epoch(self, loss_fn): self.model.train() total_loss = 0.0 for imgs, masks in self.train_loader: self.optimizer.zero_grad() preds = self.model(imgs) loss = loss_fn(preds, masks) loss.backward() self.optimizer.step() total_loss += loss.item() return total_loss / len(self.train_loader) def validate(self, loss_fn): self.model.eval() val_loss = 0.0 with torch.no_grad(): for imgs, masks in self.val_loader: preds = self.model(imgs) loss = loss_fn(preds, masks) val_loss += loss.item() self.scheduler.step(val_loss) return val_loss / len(self.val_loader)

优化器选的是 Adam，初始学习率1e-4，这两个设定在语义分割里非常稳。Adam 能自适应调整每个参数的学习率，对 U-Net 这类参数量较大的网络来说，比 SGD 更容易找到合适收敛点。ReduceLROnPlateau是学习率调度器：当验证集 loss 连续patience个 epoch 不下降时，学习率减半，模式是min，即追求验证 loss 最小化。这个调度策略尽量避免了训练后期 loss 在局部震荡出不来。模型保存的逻辑一般在每个 epoch 验证完成后：

if val_loss < best_loss: best_loss = val_loss torch.save(model.state_dict(), "models/train.pth")

只保存最优验证 loss 对应的权重，而不是每轮都存，避免磁盘空间被几十个 epoch 的权重文件占满。训练结束后models/train.pth就是最终拿来推理的权重。

4.3 main.py：入口参数与训练启动方式

main.py把以上模块组装起来，通过命令行参数控制运行模式。典型用法：

# 训练模式 python main.py --mode train --model att_unet --epochs 100 --batch_size 8 --lr 1e-4 # 评估模式 python main.py --mode eval --model att_unet --weights models/train.pth # 单图推理 python main.py --mode test --model att_unet --weights models/train.pth --input test.png

入口脚本里通常还会检查 CUDA 是否可用，自动把模型和数据搬到 GPU。batch_size的设置受限于显存，256×256 输入、batch size 8 大概需要 6 到 8GB 显存；如果显存不够，优先把 batch size 降到 4，而不是调整 patch 尺寸，因为 patch 变小会直接损失上下文信息，水体周围的建筑和阴影关系就看不清了。epochs设 100 在遥感水体任务里是够的，配合学习率调度器，通常在 40 到 60 轮就能看到验证 IoU 进入平台期。

5. 避坑排查：城市水体提取的五个高频翻车现场

遥感语义分割的坑比普通图像分割多得多，这个项目里我实际跑下来遇到了五类高频问题，按“现象 → 原因 → 解决”逐个说明。

5.1 阴影被大面积误检为水体

现象：预测结果里建筑物背光面的阴影区域被大块标成水体，IoU 虚高但视觉结果完全不可用。

原因：阴影和水体在 RGB 三个通道上的反射特性非常接近，尤其是深色建筑阴影，灰度值和水面暗部几乎重叠，仅靠颜色特征无法区分。U-Net 如果没有足够的上下文信息，很容易把阴影当水体。

解决：第一优先做训练数据增强，在预处理时把标注好的阴影区域明确标为背景，让网络见过足够多“看起来像水但实际不是”的样本。第二可以增加 NDWI 特征通道（绿波段减近红外除以两者之和），但需要多光谱数据支持。第三是训练时对 loss 做加权，把背景中阴影区域的权重调低，减少其对梯度的影响。

5.2 细长水体（河道、水渠）出现断线

现象：预测结果中河流中间断裂，原本连续的水面被切成几段，连通性差。

原因：细长水体在 U-Net 的下采样过程中空间信息被压缩，连续细长的目标在低分辨率特征图上可能只占几个像素，上采样后这些像素被淹没在背景预测里。另一个常见原因是数据集里河流样本数量少，模型没见过足够多的连续水体形态。

解决：把 patch 尺寸从 256 适当调大到 384 或 512，让单次推理覆盖更长的水体范围。如果显存不允许，可以改为重叠裁剪（stride 设为 patch 的一半），增强模型对跨 patch 连续性的感知。推理阶段还可以对相邻 patch 预测结果做加权融合，重叠区域取平均，能明显改善断线。

5.3 训练 loss 下降但验证 IoU 不涨

现象：训练集 loss 一路下降，验证集 IoU 徘徊在 60 左右上不去，分类别看水体类别 IoU 很低。

原因：这是典型的正负样本不均衡问题。城市影像里水体占比可能只有 5%，模型把所有像素预测为背景就能拿到很高的准确率，但 IoU 会很难看。另一个原因是使用了普通交叉熵损失，没有聚焦到水体区域。

解决：检查solver.py里的 lose 是否换成了 DICE 和 BCE 的组合。如果已经是组合损失，试着把dice_weight从 0.5 提高到 0.7，让优化目标更偏向区域重叠率。还可以在data_loader.py里做在线难例挖掘，按批次把预测错误率高的 patch 采样概率提高。

5.4 mask 标注与影像不对齐，训练出来边缘全是锯齿

现象：模型预测的水体边缘出现规律性锯齿，边界与真实水体轮廓偏差明显。

原因：预处理阶段对影像做了某种变换（比如旋转或裁剪），但 mask 没有执行完全相同的变换。最典型的是旋转时用了cv2.rotate处理影像，却对 mask 用了np.rot90，两个函数在非正方形影像上的行为不一致。另一个原因是影像和 mask 是从不同来源下载的，坐标系有偏移，没有做配准。

解决：把预处理脚本里的几何变换改成统一函数封装，确保影像和 mask 走同一条变换路径，这是最简单也最稳妥的做法。对新增数据，建议在训练前写一个可视化脚本，把影像和 mask 叠加画出来，人工抽检 20 张以上，确认边缘完全贴合再进训练。可视化这一步看着费时间，实际上比事后排查省事得多。

5.5 训练突然中断或 OOM 显存溢出

现象：训练跑到第 30 个 epoch 时CUDA out of memory，程序直接崩溃；或者推理时遇到单张大图直接炸显存。

原因：显存溢出多半是 batch size 和 patch 尺寸的乘积超出了显存上限，尤其在 Attention U-Net 中，多了注意力门那一层计算，比 U-Net 多占用约 20% 显存。还有一个隐蔽原因是验证集推理时没有关梯度计算，导致验证阶段显存峰值翻倍。

解决：训练脚本里在验证阶段必须包with torch.no_grad()，省下反向传播的显存占用。如果仍然溢出，降低batch_size到 4 或 2。推理大图时，不要整张图输入，用切片推理或者用torch.nn.functional.unfold分块处理，推理完再拼回原图尺寸。

6. 结果验证与推理部署：从 res.csv 到单图落地的最后一步

训练真正结束的标志不是 loss 降到最低，而是验证集指标能反映真实场景的可用性。evaluation.py和result/res.csv就是干这个活的。评估脚本通常输出 Overall Accuracy、mIoU、F1 Score、Precision、Recall 五个指标，其中 mIoU 是语义分割最通用的主指标，水体提取任务里建议额外关注 Recall——漏检的水体比误检更不可接受，因为漏掉一段河面，整个连通性分析就废了。

单图推理入口是test_one_data.py，它读取一张影像、加载models/train.pth权重、把输入缩放到训练时一致的尺寸、前向推理得到概率图，再以 0.5 为阈值转成二值掩码输出。推理代码里有一个容易被忽略的点：输入影像在预处理阶段做过的归一化参数，推理时必须用完全一样的均值和标准差，否则结果会整体偏移。我建议把归一化的 mean 和 std 写成一个公共常量文件，训练和推理都引用它，不要在两处各写一遍，这样不会出现训练好但推理结果全黑的尴尬。

对于遥感影像，一个提升可视化效果的实用技巧是输出后做形态学后处理。先用cv2.morphologyEx做开运算去掉孤立噪点，再做闭运算填补水体内部的小空洞，然后结合cv2.connectedComponents把面积小于设定阈值的连通域删掉，这能过滤掉那些零星误检。这个后处理不会大幅提升 IoU，但视觉结果会干净很多，毕设答辩时展示效果会好不少。最后把掩码和原图叠加，保存成带半透明水色的结果图，能直观证明“水体边界和道路立交桥阴影分得开”。

跑完这套流程我个人最大的感受是：水体提取这个任务，模型结构的影响远不如数据质量和后处理细节的影响大。第一次复现时我在enhance_image.py里漏了 mask 的水平翻转同步，训练出来的模型在验证集上 IoU 只有 74，修正后直接跳到 83——数据管线问题远比网络结构问题更容易吃掉分数。从那以后我每次做遥感分割项目，都会强制在训练前跑一遍“原图 + mask 叠加可视化”的抽检脚本，确认数据没问题才敢开训练。这条习惯帮我省掉过太多半夜排查的精力，希望也能帮到你。

本文还有配套的精品资源，点击获取