PyTorch U-Net城市遥感水体识别实战：从假彩色影像到二值掩膜全流程-开发者社区

简介：这是一套面向城市区域水体提取的开箱即用实战资源，适合遥感、计算机视觉方向的学生与开发者，可用于本科毕业设计、课程大作业或轻量级水体监测原型开发。资源以高分辨率城市遥感影像为对象，集成U-Net与AttU-Net两种主流分割网络，覆盖从数据预处理、加载、模型定义到训练、评估与单图推理的完整链路，并附带已训练权重与典型可视化结果，便于快速验证与二次开发。压缩包共32个文件，约733KB，以13个png可视化图、11个py脚本、3个csv标签文件为主，另含1个pth权重、1个md说明文档及配置文件，结构清晰、注释详尽。目前已有134人学习下载。读者可直接获得可复现的训练推理代码、Urban_pre子集与CSV标签、增强图像脚本及结果导出工具，无需复杂配置即可跑通全流程，并借助README快速理解各模块职责与排错思路。

1. 城市遥感图水体识别：从一张假彩色影像到可交付的掩膜

手里有一景城市遥感影像，想快速把河道、湖泊、坑塘的水面边界抠出来，靠人眼描边一下午也画不完一景，换个季节的影像又得重来。这个实战包解决的正是这件事：它把城市遥感场景下的水体提取拆成了一条能跑通的链路——PyTorch 实现的 U-Net 分割模型、配套的标注数据集、训练与推理脚本，以及从原始影像到二值掩膜的完整流程。适合做遥感图像分割的算法同学、做城市水务或规划相关数据处理的一线工程师，以及想拿一个真实分割任务练手 PyTorch 的人。它不追求 SOTA，追求的是你拿到手能复现、能改、能换自己的数据继续训。下面按“这东西怎么搭起来、参数怎么设、哪里容易翻车”的顺序讲。

2. 水体提取为什么选 U-Net：分割任务与数据管线的选型逻辑

2.1 城市遥感水体提取的任务特性

城市遥感里的水体提取，本质是一个二分类语义分割问题：每个像素判成“水”或“非水”。听起来简单，但城市场景有几个绕不开的麻烦。第一，水体形态极不规整，有宽阔江面，也有细如发丝的排水渠，还有被桥梁、船只、树冠遮挡后断裂的水面。第二，城市地物里有一堆“假水”：深色沥青路面、建筑阴影、蓝色屋顶、光伏板，它们在假彩色或真彩色影像上和水的光谱响应很接近，模型很容易把它们误判成水。第三，遥感影像的尺度差异大，同一景里既有几百米宽的湖，也有几米宽的沟渠，单一尺度的卷积核很难同时兼顾。

这些特性决定了模型必须能融合多尺度特征，并且对边界敏感。U-Net 的编码器-解码器结构加跳跃连接，恰好把浅层的高分辨率细节和深层的语义信息拼在一起，对小面积水体和细长水系的边界保留比较友好。这也是它在遥感水体提取里长期被当作 baseline 的原因——不是因为它最强，而是因为它稳定、易训、改起来直观。

2.2 数据集组织与标注格式

实战包里的数据集一般按影像-掩膜成对组织，目录结构常见做法是这样：

dataset/ ├── images/ │ ├── tile_0001.png │ ├── tile_0002.png │ └── ... ├── masks/ │ ├── tile_0001.png │ ├── tile_0002.png │ └── ... └── splits/ ├── train.txt ├── val.txt └── test.txt

影像可以是三通道真彩色，也可以是包含近红外波段的四通道假彩色。掩膜是单通道灰度图，水体像素值 255，背景 0。splits 里的 txt 每行一个样本名，不带扩展名，训练脚本按行读取再拼路径。这种组织方式的好处是切分逻辑和数据解耦，换数据集时只改 txt 就行。

提示：掩膜务必确认是单通道且只有 0/255 两个值。我见过有人拿 RGB 三通道掩膜直接喂进去，模型照样能训，但 loss 一直下不去，排查半天才发现是通道数不对。

2.3 数据增强与归一化参数

遥感影像的增强不能照搬自然图像那套。水平翻转、垂直翻转、90 度旋转是安全的，因为遥感视角下地物方向没有语义偏好。但颜色抖动要慎用，水体和非水体的区分很大程度依赖光谱，色相偏移过大反而制造噪声。归一化建议按数据集统计的均值和标准差来，而不是直接用 ImageNet 的默认值，因为遥感影像的分布和自然图像差得远。

import numpy as np from PIL import Image def compute_mean_std(image_paths, channels=3): # 逐通道累加像素值，用于统计数据集自身的均值和标准差 sums = np.zeros(channels) sq_sums = np.zeros(channels) count = 0 for p in image_paths: img = np.asarray(Image.open(p)).astype(np.float32) / 255.0 if img.ndim == 2: img = img[:, :, None] sums += img.reshape(-1, channels).sum(axis=0) sq_sums += (img.reshape(-1, channels) ** 2).sum(axis=0) count += img.shape[0] * img.shape[1] mean = sums / count std = np.sqrt(sq_sums / count - mean ** 2) return mean, std

这段代码遍历所有训练影像，按通道统计均值和标准差。channels 要和实际输入通道数一致，四通道假彩色就传 4。算出来的 mean 和 std 写进配置，训练和推理必须用同一组值，否则推理时输入分布和训练不一致，输出会明显变差。这是很多人忽略的坑：训练时归一化了，推理时忘了，结果模型“突然不会了”。

3. U-Net 模型搭建与训练脚本：从网络结构到 loss 选择

3.1 编码器-解码器结构与通道配置

实战包里的 U-Net 通常是四层下采样、四层上采样的经典结构。编码器每层两个 3x3 卷积加 ReLU，然后 2x2 最大池化；解码器每层先上采样，再和对应编码层的特征拼接，然后两个 3x3 卷积。通道数从 64 开始逐层翻倍到 512，再逐层减回 64，最后 1x1 卷积输出单通道 logits。

import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.block = nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), ) def forward(self, x): return self.block(x) class UNet(nn.Module): def __init__(self, in_ch=3, out_ch=1, base=64): super().__init__() # 编码器：通道数 base -> base*8 self.enc1 = DoubleConv(in_ch, base) self.enc2 = DoubleConv(base, base * 2) self.enc3 = DoubleConv(base * 2, base * 4) self.enc4 = DoubleConv(base * 4, base * 8) self.pool = nn.MaxPool2d(2) # 瓶颈层 self.bottleneck = DoubleConv(base * 8, base * 16) # 解码器：上采样后拼接，通道数减半 self.up4 = nn.ConvTranspose2d(base * 16, base * 8, 2, stride=2) self.dec4 = DoubleConv(base * 16, base * 8) self.up3 = nn.ConvTranspose2d(base * 8, base * 4, 2, stride=2) self.dec3 = DoubleConv(base * 8, base * 4) self.up2 = nn.ConvTranspose2d(base * 4, base * 2, 2, stride=2) self.dec2 = DoubleConv(base * 4, base * 2) self.up1 = nn.ConvTranspose2d(base * 2, base, 2, stride=2) self.dec1 = DoubleConv(base * 2, base) self.head = nn.Conv2d(base, out_ch, 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)) b = self.bottleneck(self.pool(e4)) d4 = self.dec4(torch.cat([self.up4(b), e4], dim=1)) d3 = self.dec3(torch.cat([self.up3(d4), e3], dim=1)) d2 = self.dec2(torch.cat([self.up2(d3), e2], dim=1)) d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1)) return self.head(d1)

base 控制模型宽度，默认 64。显存吃紧就降到 32，精度会掉一点但能跑起来。in_ch 按影像通道数改，out_ch 二分类固定为 1。拼接时用 torch.cat 沿通道维，注意上采样后的尺寸要和对应编码层一致，输入尺寸不是 16 的整数倍时会出现尺寸对不齐，常见做法是训练前把影像裁成 16 的倍数。

3.2 损失函数与类别不平衡处理

水体在整景影像里占比往往不高，背景像素远多于水体，直接用 BCE 会让模型倾向于全预测背景。实战包里一般用 BCE 和 Dice 的组合损失，Dice 对前景占比低的情况更鲁棒。

class BCEDiceLoss(nn.Module): def __init__(self, bce_weight=0.5): super().__init__() self.bce = nn.BCEWithLogitsLoss() self.bce_weight = bce_weight def forward(self, logits, targets): bce_loss = self.bce(logits, targets) probs = torch.sigmoid(logits) # 展平后计算 Dice，避免空间维度影响 probs = probs.view(probs.size(0), -1) targets = targets.view(targets.size(0), -1) intersection = (probs * targets).sum(dim=1) dice = (2 * intersection + 1e-6) / (probs.sum(dim=1) + targets.sum(dim=1) + 1e-6) dice_loss = 1 - dice.mean() return self.bce_weight * bce_loss + (1 - self.bce_weight) * dice_loss

bce_weight 默认 0.5，水体占比极低时可以调到 0.3，让 Dice 主导。平滑项 1e-6 防止除零。这个组合的好处是 BCE 提供稳定的梯度，Dice 直接优化重叠度，两者互补。如果训练时 loss 震荡厉害，先检查 targets 是不是 float 类型且值在 0 到 1 之间，整型会直接报错或静默出错。

3.3 训练循环与关键超参

训练脚本的核心是数据加载、前向、loss 反传、优化器更新这几步。优化器常用 Adam，学习率 1e-3 起步，配合余弦退火或 ReduceLROnPlateau。batch size 受显存限制，常见 4 到 8。

from torch.utils.data import DataLoader from dataset import WaterDataset # 实战包里的数据集类 def train_one_epoch(model, loader, optimizer, criterion, device): model.train() total_loss = 0.0 for imgs, masks in loader: imgs = imgs.to(device) masks = masks.to(device).float() optimizer.zero_grad() logits = model(imgs) loss = criterion(logits, masks) loss.backward() optimizer.step() total_loss += loss.item() * imgs.size(0) return total_loss / len(loader.dataset)

device 优先 cuda，没有就 cpu 但会很慢。masks 转 float 是因为 BCEWithLogitsLoss 要求目标为浮点。每个 epoch 后在验证集上算 IoU 和 Dice，保存指标最好的权重，别只看 loss，loss 低不代表掩膜质量好。学习率调度建议用 ReduceLROnPlateau，patience 设 5，验证 IoU 连续 5 个 epoch 不升就降一半。

4. 推理与后处理：把 logits 变成能用的水体掩膜

4.1 滑窗推理与大图拼接

遥感影像往往很大，整景塞不进显存，常见做法是滑窗切块推理再拼回去。窗口大小和训练时一致，比如 512x512，步长设成窗口的一半做重叠，重叠区域取平均或取最大值来消除拼接缝。

import torch import numpy as np def sliding_inference(model, image, window=512, stride=256, device="cuda"): model.eval() h, w = image.shape[:2] prob_map = np.zeros((h, w), dtype=np.float32) count_map = np.zeros((h, w), dtype=np.float32) with torch.no_grad(): for y in range(0, h, stride): for x in range(0, w, stride): y2 = min(y + window, h) x2 = min(x + window, w) y1 = max(0, y2 - window) x1 = max(0, x2 - window) patch = image[y1:y2, x1:x2] tensor = torch.from_numpy(patch).permute(2, 0, 1).float().unsqueeze(0).to(device) logits = model(tensor) prob = torch.sigmoid(logits).squeeze().cpu().numpy() prob_map[y1:y2, x1:x2] += prob count_map[y1:y2, x1:x2] += 1 return prob_map / np.maximum(count_map, 1e-6)

stride 越小重叠越多，拼接越平滑但推理越慢。window 必须和训练时的输入尺寸匹配，否则 BatchNorm 的统计会偏。count_map 记录每个像素被覆盖的次数，最后取平均，避免边缘区域因为覆盖少而偏暗。这段是推理里最容易出问题的地方，拼接缝、边缘缺失基本都出在这。

4.2 阈值选取与形态学后处理

概率图出来后要二值化，阈值默认 0.5，但实际场景里可以调。水体召回不够就降到 0.4，误检多就升到 0.6。二值化后常有零星噪点和孔洞，用形态学开运算去噪、闭运算填洞。

import cv2 def postprocess(prob_map, thresh=0.5, min_area=50): binary = (prob_map > thresh).astype(np.uint8) * 255 kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)) binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel) binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel) # 去掉面积过小的连通域，抑制噪点 num, labels, stats, _ = cv2.connectedComponentsWithStats(binary, connectivity=8) for i in range(1, num): if stats[i, cv2.CC_STAT_AREA] < min_area: binary[labels == i] = 0 return binary

min_area 按影像分辨率调，高分辨率影像可以设大一点，比如 100。开运算去孤立噪点，闭运算补小孔洞，顺序别反。连通域过滤是最后一道防线，把面积过小的误检清掉。这套后处理不复杂，但能明显改善目视效果。

5. 避坑与排查：水体提取里最容易翻车的五件事

5.1 掩膜值不是 0/255 导致 loss 异常

现象：训练 loss 一开始就很大，或者一直不降。原因：掩膜是 0/1 而不是 0/255，或者掩膜是 RGB 三通道。解决：训练前统一检查掩膜，用脚本统计唯一值，确保是单通道且只有 0 和 255，读入后除以 255 转成 0/1 浮点。

5.2 训练和推理归一化不一致

现象：训练时验证 IoU 正常，推理自己的影像结果一塌糊涂。原因：推理时忘了用训练统计的 mean/std，或者用了 ImageNet 默认值。解决：把 mean/std 写进配置文件，训练和推理读同一份，推理脚本里显式打印实际用的归一化参数核对。

5.3 输入尺寸不是 16 的倍数导致拼接报错

现象：前向传播时报尺寸不匹配，torch.cat 失败。原因：U-Net 四次下采样，输入宽高必须是 16 的整数倍，否则上采样后和编码层尺寸对不上。解决：训练和推理前把影像 padding 或裁切到 16 的倍数，推理时记录原始尺寸，输出后再裁回去。

5.4 深色地物被误判成水体

现象：掩膜里出现大片路面、阴影或蓝色屋顶。原因：这些地物光谱和水的相似度高，模型没学够区分特征。解决：在训练集里补充这类负样本，或者加一个“难例挖掘”环节，把误检多的区域裁出来重新标注加入训练。也可以在 loss 里对假阳性加权重。

5.5 滑窗推理拼接缝明显

现象：输出掩膜上有网格状接缝。原因：stride 等于 window，没有重叠，或者重叠区域直接覆盖而不是平均。解决：stride 设成 window 的一半，重叠区域用概率平均，别用最后一次覆盖的结果。如果还有缝，检查 count_map 是否在边界处为 0 导致除零。

6. 进阶技巧：用测试时增强和阈值搜索把 IoU 再抬一截

模型训完不是终点，推理阶段还有免费的精度可以捡。测试时增强（TTA）是性价比最高的一招：对同一张影像做水平翻转、垂直翻转、90 度旋转，分别推理后再把结果逆变换回来取平均。水体在遥感视角下没有方向偏好，这些变换不会引入语义错误，但能让预测更稳。实测里 TTA 通常能把 IoU 抬 1 到 3 个百分点，代价是推理时间翻几倍，看你能不能接受。

def tta_inference(model, image, device="cuda"): # 四种变换：原图、水平翻转、垂直翻转、旋转180度 variants = [ image, np.fliplr(image), np.flipud(image), np.rot90(image, 2), ] probs = [] for v in variants: tensor = torch.from_numpy(np.ascontiguousarray(v)).permute(2, 0, 1).float().unsqueeze(0).to(device) with torch.no_grad(): p = torch.sigmoid(model(tensor)).squeeze().cpu().numpy() probs.append(p) # 逆变换后平均 probs[1] = np.fliplr(probs[1]) probs[2] = np.flipud(probs[2]) probs[3] = np.rot90(probs[3], -2) return np.mean(probs, axis=0)

逆变换的顺序要和正变换相反，rot90 用 -2 转回来。这段代码假设输入是 HWC 格式，如果推理管线里是 CHW 要相应调整。

另一个技巧是阈值搜索。0.5 不是万能值，在验证集上扫一遍阈值，从 0.3 到 0.7 每隔 0.05 算一次 IoU，取最高的那个作为最终阈值。这个操作几乎零成本，但经常能比默认 0.5 好一截。我一般会把阈值搜索和 TTA 一起做，先 TTA 得到概率图，再在验证集上搜阈值，最后固化成推理配置。