基于开源遥感水体分割数据集，从零构建U-Net模型实战指南

简介：语义分割是计算机视觉的核心任务之一，旨在为图像中的每个像素分配类别标签，其原理是通过深度学习模型学习像素与语义类别之间的映射关系。这项技术在遥感图像解译领域具有极高的技术价值，是实现自动化、精细化地物识别与分析的关键。在众多应用场景中，水体分割对于洪涝监测、水资源管理等至关重要。本文聚焦于一个开箱即用的遥感水体分割数据集，详细阐述了如何利用PyTorch框架和U-Net模型，从环境配置、数据加载、模型训练到调优部署，完成一个完整的语义分割项目实战。内容涵盖了数据增强、损失函数选择、后处理技巧以及针对遥感影像特性的优化策略，为相关研究和工程应用提供了清晰的实践路径。

1. 项目概述：一份“开箱即用”的遥感水体分割数据集

如果你正在研究遥感图像分析，特别是水体识别与提取，那么你很可能遇到过数据获取和处理的巨大障碍。从卫星影像下载、筛选、标注到格式统一，每一步都耗时费力，而且对计算资源有不低的要求。今天要聊的这个项目，可以说是一个“及时雨”——一个名为“卫星拍摄下的水体图像语义分割数据集”的资源包。它包含了约2300张已经处理完毕的卫星图像和对应的像素级标签，专门用于二分类（水体 vs. 非水体）的语义分割任务，拿到手就能直接扔进模型里开始训练。

这个数据集的核心价值在于它的“即用性”。在遥感领域，数据是模型性能的基石，但高质量、大规模、标注精准的公开数据集却相对稀缺。很多研究者，尤其是学生或独立开发者，宝贵的时间和算力往往被卡在数据预处理阶段。这个数据集直接解决了这个问题，它提供的是已经完成裁剪、配准、标注和格式转换的“成品”，让你能跳过最繁琐的步骤，直接聚焦于模型架构设计、调参和性能优化。对于想快速验证一个水体分割算法想法，或是作为课程设计、毕业项目的数据基础，它都是一个非常理想的起点。

2. 数据集深度解析：从像素到应用场景

2.1 数据来源与特性拆解

首先，我们来拆解一下这个数据集可能包含的信息。标题中“卫星拍摄下”是一个关键描述，这意味着影像数据源来自各类对地观测卫星。常见的开源数据源可能包括：

• 哨兵系列卫星：如Sentinel-2，提供10米、20米、60米等多光谱数据，是地表监测的主力，其影像免费且更新频率高。
• Landsat系列卫星：如Landsat 8/9，提供30米分辨率的多光谱数据，历史存档丰富，适合长时间序列分析。
• 高分系列卫星：部分数据可能来源于国内外的高分辨率对地观测系统。

数据集约2300张的规模，在遥感分割任务中属于中等偏上的体量，足以支撑一个中等复杂度模型（如U-Net、DeepLabv3+）的有效训练，避免因数据量不足导致的过拟合。所谓“已处理完”，通常意味着以下工作已经完成：

1. 影像筛选与裁剪：从大范围的原始卫星影像中，截取出包含丰富水体特征（如河流、湖泊、水库、近海）的区域。
2. 辐射定标与大气校正：将卫星的原始数字量化值转换为地表反射率，消除大气散射、吸收的影响，使不同时间、不同传感器获取的影像具有可比性。这是遥感分析的基础，如果没做，模型学到的可能是“大气噪声”而非地物特征。
3. 多波段合成与选择：卫星影像通常包含多个波段（如红、绿、蓝、近红外、短波红外）。水体在近红外和短波红外波段吸收强、反射率极低，与植被、土壤形成鲜明对比。数据集很可能已经选择了最适合水体识别的波段组合（例如，使用归一化水体指数NDWI所需的绿光和近红外波段）。
4. 标签制作：这是最核心、最耗时的一步。“2类别图像分割”意味着每张影像都配有一张同样大小的单通道掩膜图像。在这个掩膜中，每个像素被标记为两类之一：水体（通常用像素值1或255表示）和非水体（用0表示）。标签的精度直接决定了模型性能的上限。

2.2 核心应用场景与价值

这个数据集的应用场景非常广泛，远不止于一个单纯的算法练习。它能够直接支撑起一系列具有实际价值的应用：

• 洪涝灾害监测与评估：训练好的模型可以快速从灾后卫星影像中提取淹没区范围，为救灾决策和损失评估提供关键信息。
• 水资源调查与管理：用于定期监测湖泊、水库的水面面积变化，评估蓄水量，或统计河流网络的分布与变迁。
• 生态环境研究：追踪湿地范围变化、海岸线侵蚀或增长，研究水体与周边生态系统的互动关系。
• 城市规划与基础建设：辅助进行城市水域规划，监测人工水体（如景观湖、污水处理厂）的建设情况。
• 农业灌溉分析：识别农田中的灌溉区域和沟渠分布。

它的价值在于提供了一个标准化的“试验场”。不同研究者可以使用同一套数据来公平地比较不同分割模型（如FCN, U-Net, PSPNet, SegFormer）的性能，推动了算法研究的可比性和复现性。同时，对于初学者而言，它降低了进入遥感AI领域的门槛，让大家能把精力集中在理解模型原理和调参技巧上。

3. 数据准备与模型训练环境搭建

拿到一个即用数据集后，第一步不是急着跑代码，而是先“认识”它，并搭建好实验环境。

3.1 数据集结构检查与理解

通常，一个组织良好的分割数据集文件夹结构如下：

WaterBody_Seg_Dataset/ ├── images/ # 存放所有卫星影像图片 │ ├── 0001.png │ ├── 0002.png │ └── ... ├── masks/ # 存放所有对应的标签掩膜图片 │ ├── 0001.png │ ├── 0002.png │ └── ... └── README.md # 数据说明文件（非常重要！）

你需要首先检查这个结构。更重要的是，仔细阅读README.md文件。它应该包含以下关键信息：

• 影像规格：图像尺寸（如256x256, 512x512）、通道数（是RGB三通道还是包含近红外的四通道？）、数据类型（uint8, float32？）。
• 标签编码：明确水体和非水体像素的具体数值（例如，0代表背景，1代表水体；或者0背景，255水体）。
• 数据划分建议：是否提供了训练集/验证集/测试集的划分列表？如果没有，你需要自己按比例（如7:2:1）随机划分，并确保划分是可复现的（固定随机种子）。
• 来源与许可：说明数据来源，以及使用许可（如CC-BY协议），确保你的后续使用是合规的。

一个快速检查数据匹配和标签可视化的小脚本至关重要：

import os import cv2 import numpy as np import matplotlib.pyplot as plt image_path = './images/0001.png' mask_path = './masks/0001.png' img = cv2.imread(image_path) mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) # 以灰度图方式读取标签 print(f"图像形状: {img.shape}, 数据类型: {img.dtype}") print(f"标签形状: {mask.shape}, 数据类型: {mask.dtype}") print(f"标签唯一值: {np.unique(mask)}") # 确认标签编码 # 可视化 fig, axes = plt.subplots(1, 3, figsize=(12,4)) axes[0].imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) axes[0].set_title('卫星影像') axes[0].axis('off') axes[1].imshow(mask, cmap='gray') axes[1].set_title('标签掩膜') axes[1].axis('off') # 叠加显示 axes[2].imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) axes[2].imshow(mask, cmap='jet', alpha=0.5) # 半透明叠加 axes[2].set_title('影像-标签叠加') axes[2].axis('off') plt.show()

3.2 深度学习环境配置要点

对于语义分割任务，主流的框架如PyTorch或TensorFlow/Keras都是不错的选择。这里以PyTorch为例，环境配置的核心是版本兼容性。

# 使用conda创建虚拟环境是最佳实践，避免包冲突 conda create -n rs_seg python=3.8 -y conda activate rs_seg # 安装PyTorch（请根据你的CUDA版本去官网获取对应命令） # 例如，对于CUDA 11.3 pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113 # 安装常用工具库 pip install opencv-python pillow matplotlib scikit-learn scikit-image tqdm tensorboard # 安装一个常用的分割模型库，如segmentation-models-pytorch (smp) pip install segmentation-models-pytorch

注意：segmentation-models-pytorch库封装了许多预编码的经典分割模型（U-Net, DeepLabV3+, FPN等）和预训练骨干网络（ResNet, EfficientNet等），能极大提升开发效率。但要注意其与PyTorch版本的兼容性。

4. 模型选择、训练策略与核心代码实现

有了数据和环境，接下来就是选择模型并开始训练。对于2300张二分类水体分割任务，我们不需要一开始就追求最复杂的模型，应从经典且高效的模型入手。

4.1 模型选择与构建

U-Net以其编码器-解码器结构和跳跃连接，在医学影像和遥感分割中久经考验，特别适合中等规模数据集。我们可以利用smp库快速构建一个U-Net：

import torch import segmentation_models_pytorch as smp import torch.nn as nn # 定义模型 model = smp.Unet( encoder_name="resnet34", # 编码器骨干网络，resnet34在速度和精度间取得平衡 encoder_weights="imagenet", # 使用在ImageNet上预训练的权重，加速收敛 in_channels=3, # 输入通道数，根据你的数据调整（RGB为3） classes=2, # 分类数：2类（水体、非水体） ) # 将模型移动到GPU device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = model.to(device) # 打印模型概要 print(model)

这里选择ResNet34作为编码器，是因为它比ResNet18特征提取能力更强，又比ResNet50/101更轻量，适合作为基线模型。in_channels需要根据你的数据确认，如果是RGB就是3，如果包含了近红外波段可能就是4。

4.2 数据加载与预处理流程

高效的数据管道是训练稳定的关键。我们需要自定义一个Dataset类来读取图像-标签对，并施加必要的增强。

from torch.utils.data import Dataset, DataLoader import albumentations as A from albumentations.pytorch import ToTensorV2 class WaterBodyDataset(Dataset): def __init__(self, image_dir, mask_dir, transform=None): self.image_dir = image_dir self.mask_dir = mask_dir self.transform = transform self.images = sorted(os.listdir(image_dir)) self.masks = sorted(os.listdir(mask_dir)) # 简易检查：确保图像和标签文件名能对应上 assert len(self.images) == len(self.masks), "图像和标签数量不匹配！" for img, msk in zip(self.images, self.masks): assert os.path.splitext(img)[0] == os.path.splitext(msk)[0], f"文件名不匹配: {img} vs {msk}" def __len__(self): return len(self.images) def __getitem__(self, idx): img_path = os.path.join(self.image_dir, self.images[idx]) mask_path = os.path.join(self.mask_dir, self.masks[idx]) # 使用OpenCV读取，注意通道顺序和标签格式 image = cv2.imread(img_path) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # 转为RGB mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) # 重要：将标签归一化到0-1范围，或者处理为0/1二值 # 假设原始标签中水体为255，非水体为0 mask = (mask > 128).astype(np.uint8) # 阈值化，确保标签是0和1 if self.transform: augmented = self.transform(image=image, mask=mask) image = augmented['image'] mask = augmented['mask'] # 将mask从(H, W)转为(H, W, 1)以便于某些损失函数计算，但通常Pytorch期望是LongTensor mask = mask.long() if torch.is_tensor(mask) else torch.from_numpy(mask).long() return image, mask # 定义训练和验证的数据增强 # 训练时使用更强的增强以提升模型鲁棒性 train_transform = A.Compose([ A.RandomRotate90(p=0.5), A.HorizontalFlip(p=0.5), A.VerticalFlip(p=0.5), A.RandomBrightnessContrast(brightness_limit=0.1, contrast_limit=0.1, p=0.3), # 遥感影像增强慎用色彩抖动，可能破坏光谱特性 A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), # ImageNet归一化参数，适用于预训练骨干 ToTensorV2(), ]) # 验证时只需归一化和Tensor转换 val_transform = A.Compose([ A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), ToTensorV2(), ]) # 创建数据集和数据加载器 train_dataset = WaterBodyDataset(train_img_dir, train_mask_dir, transform=train_transform) val_dataset = WaterBodyDataset(val_img_dir, val_mask_dir, transform=val_transform) train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, num_workers=4, pin_memory=True) val_loader = DataLoader(val_dataset, batch_size=8, shuffle=False, num_workers=4, pin_memory=True)

实操心得：数据增强是提升模型泛化能力的利器，但对于遥感影像，增强策略需要谨慎。几何变换（旋转、翻转）通常是安全的，但色彩空间变换（亮度、对比度、饱和度）可能会改变地物的光谱反射特性，需要根据任务测试效果。我个人的经验是，对于水体分割，轻微的亮度对比度调整可能有益（模拟不同光照条件），但大幅度的色彩抖动可能有害。

4.3 损失函数、评估指标与训练循环

语义分割常用的损失函数是Dice Loss或交叉熵损失（CrossEntropy Loss），也可以结合使用。评估指标常用交并比（IoU）和Dice系数。

import torch.optim as optim from torch.nn import CrossEntropyLoss from sklearn.metrics import jaccard_score def dice_coeff(pred, target, smooth=1e-6): # pred: (N, C, H, W) after softmax # target: (N, H, W) with class indices pred = torch.argmax(pred, dim=1) # 将预测结果转为类别索引 pred_flat = pred.contiguous().view(-1) target_flat = target.contiguous().view(-1) intersection = (pred_flat * target_flat).sum() dice = (2. * intersection + smooth) / (pred_flat.sum() + target_flat.sum() + smooth) return dice.item() # 初始化 criterion = CrossEntropyLoss() # 内部会处理类别不平衡（需结合权重或采样） optimizer = optim.Adam(model.parameters(), lr=1e-4) scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=5, verbose=True) num_epochs = 50 best_iou = 0.0 for epoch in range(num_epochs): model.train() train_loss = 0.0 train_dice = 0.0 for images, masks in train_loader: images, masks = images.to(device), masks.to(device) optimizer.zero_grad() outputs = model(images) # (N, 2, H, W) loss = criterion(outputs, masks) loss.backward() optimizer.step() train_loss += loss.item() with torch.no_grad(): train_dice += dice_coeff(outputs, masks) avg_train_loss = train_loss / len(train_loader) avg_train_dice = train_dice / len(train_loader) # 验证阶段 model.eval() val_loss = 0.0 val_dice = 0.0 all_preds = [] all_targets = [] with torch.no_grad(): for images, masks in val_loader: images, masks = images.to(device), masks.to(device) outputs = model(images) loss = criterion(outputs, masks) val_loss += loss.item() val_dice += dice_coeff(outputs, masks) # 为计算IoU收集数据 preds = torch.argmax(outputs, dim=1).cpu().numpy().flatten() targets = masks.cpu().numpy().flatten() all_preds.extend(preds) all_targets.extend(targets) avg_val_loss = val_loss / len(val_loader) avg_val_dice = val_dice / len(val_loader) val_iou = jaccard_score(all_targets, all_preds, average='binary') # 二分类IoU print(f'Epoch [{epoch+1}/{num_epochs}]') print(f' Train Loss: {avg_train_loss:.4f}, Train Dice: {avg_train_dice:.4f}') print(f' Val Loss: {avg_val_loss:.4f}, Val Dice: {avg_val_dice:.4f}, Val IoU: {val_iou:.4f}') # 学习率调整 scheduler.step(val_iou) # 保存最佳模型 if val_iou > best_iou: best_iou = val_iou torch.save({ 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'best_iou': best_iou, }, 'best_model.pth') print(f' -> 保存最佳模型，IoU: {best_iou:.4f}')

注意事项：类别不平衡是水体分割的常见问题。一张图中水体像素可能远少于非水体像素。单纯使用CrossEntropyLoss可能导致模型偏向背景。解决方法有：1) 在损失函数中为水体类别设置更高的权重（weight参数）；2) 使用DiceLoss或FocalLoss这类本身能处理不平衡问题的损失函数；3) 在数据加载时对包含水体的样本进行过采样。建议先尝试设置类别权重，权重可以粗略地设为“非水体像素数 / 水体像素数”。

5. 训练优化技巧与超参数调优实战

模型跑起来只是第一步，要想获得好的分割效果，调优过程必不可少。这里分享几个针对遥感水体分割的关键调优点。

5.1 学习率策略与优化器选择

学习率是训练中最重要的超参数之一。对于使用预训练编码器的模型，通常采用分层学习率策略：编码器（特征提取部分）使用较小的学习率进行微调，解码器（新添加的部分）使用较大的学习率。

# 分层设置学习率示例 encoder_params = [] decoder_params = [] for name, param in model.named_parameters(): if 'encoder' in name: encoder_params.append(param) else: decoder_params.append(param) optimizer = optim.Adam([ {'params': encoder_params, 'lr': 1e-5}, # 编码器小学习率微调 {'params': decoder_params, 'lr': 1e-4} # 解码器较大学习率快速学习 ]) # 配合余弦退火热身（CosineAnnealingWarmRestarts）学习率调度器 scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=2, eta_min=1e-6)

Adam优化器通常是默认的可靠选择。也可以尝试AdamW（Adam with weight decay），它往往有更好的泛化性能。初始学习率设置在1e-4附近进行搜索是一个好的起点。

5.2 针对遥感影像的数据增强特化

通用的图像增强可能不适用于遥感数据。我们需要设计对光谱和空间特征有意义的增强：

import albumentations as A specialized_transform = A.Compose([ # 几何增强：对遥感影像非常安全有效 A.RandomRotate90(p=0.5), A.HorizontalFlip(p=0.5), A.VerticalFlip(p=0.5), A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.1, rotate_limit=45, p=0.5, border_mode=cv2.BORDER_CONSTANT, value=0), # 光学增强：模拟不同大气、光照条件，但幅度要小 A.RandomBrightnessContrast(brightness_limit=0.15, contrast_limit=0.15, p=0.3), # 模拟轻微云层遮挡或噪声（谨慎使用） A.CoarseDropout(max_holes=8, max_height=16, max_width=16, fill_value=0, p=0.2), # 归一化 A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), ToTensorV2(), ])

CoarseDropout（随机遮挡）可以模拟云、云影或传感器瑕疵，提升模型对局部遮挡的鲁棒性，但填充值（fill_value）需要谨慎设置，对于归一化后的影像，用0或均值填充是常见选择。

5.3 后处理与模型集成提升精度

模型直接输出的预测图往往存在一些小的空洞或孤立的噪声点。简单的后处理能显著提升视觉效果和定量指标。

import cv2 import numpy as np from scipy import ndimage def postprocess_mask(pred_mask, area_threshold=50): """ pred_mask: 二值化后的预测图 (0, 1) """ # 1. 连通域分析，去除小面积噪声 num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(pred_mask.astype(np.uint8), connectivity=8) processed_mask = np.zeros_like(pred_mask) for i in range(1, num_labels): # 跳过背景标签0 if stats[i, cv2.CC_STAT_AREA] >= area_threshold: processed_mask[labels == i] = 1 # 2. 可选：形态学闭运算，填充细小空洞，平滑边界 kernel = np.ones((3,3), np.uint8) processed_mask = cv2.morphologyEx(processed_mask, cv2.MORPH_CLOSE, kernel, iterations=1) return processed_mask

对于追求极致精度的场景，可以尝试模型集成。最简单的方法是测试时增强：对同一张输入图像进行多种变换（如原图、水平翻转、垂直翻转），分别预测，然后将结果反变换回来进行平均或投票。

def tta_inference(model, image, device): """ 简单的测试时增强：原图 + 水平翻转 """ model.eval() with torch.no_grad(): # 原图预测 image_tensor = val_transform(image=image)['image'].unsqueeze(0).to(device) pred1 = torch.softmax(model(image_tensor), dim=1)[:, 1, ...].cpu().numpy() # 取水体类别的概率 # 水平翻转后预测 image_flipped = cv2.flip(image, 1) image_tensor_flipped = val_transform(image=image_flipped)['image'].unsqueeze(0).to(device) pred2_flipped = torch.softmax(model(image_tensor_flipped), dim=1)[:, 1, ...].cpu().numpy() pred2 = cv2.flip(pred2_flipped.squeeze(), 1) # 翻转回来 # 平均概率 avg_pred = (pred1.squeeze() + pred2) / 2.0 final_mask = (avg_pred > 0.5).astype(np.uint8) # 阈值化 return final_mask

6. 结果分析、可视化与常见问题排查

训练完成后，我们需要系统地评估模型表现，并学会诊断问题。

6.1 定量与定性分析模型性能

除了在验证集上计算整体的IoU和Dice，更重要的是进行按样本分析和错误案例分析。

def evaluate_and_visualize(model, dataloader, device, save_dir='results'): model.eval() os.makedirs(save_dir, exist_ok=True) iou_list = [] with torch.no_grad(): for idx, (images, masks) in enumerate(dataloader): images, masks = images.to(device), masks.to(device) outputs = model(images) preds = torch.argmax(outputs, dim=1).cpu().numpy() # (N, H, W) targets = masks.cpu().numpy() for i in range(images.size(0)): pred = preds[i] target = targets[i] iou = jaccard_score(target.flatten(), pred.flatten(), average='binary') iou_list.append(iou) # 保存可视化结果对比 if idx < 5: # 只保存前几个批次 img_np = images[i].cpu().permute(1,2,0).numpy() # 反归一化用于显示 mean = np.array([0.485, 0.456, 0.406]) std = np.array([0.229, 0.224, 0.225]) img_np = std * img_np + mean img_np = np.clip(img_np, 0, 1) fig, axes = plt.subplots(1, 4, figsize=(16,4)) axes[0].imshow(img_np) axes[0].set_title('Input Image') axes[0].axis('off') axes[1].imshow(target, cmap='gray') axes[1].set_title(f'Ground Truth') axes[1].axis('off') axes[2].imshow(pred, cmap='gray') axes[2].set_title(f'Prediction (IoU: {iou:.3f})') axes[2].axis('off') # 显示错误区域（红色为假阳性，蓝色为假阴性） error_map = np.zeros((*pred.shape, 3)) error_map[(pred==1) & (target==0)] = [1,0,0] # 假阳性，红色 error_map[(pred==0) & (target==1)] = [0,0,1] # 假阴性，蓝色 axes[3].imshow(img_np) axes[3].imshow(error_map, alpha=0.6) axes[3].set_title('Error Overlay (FP:Red, FN:Blue)') axes[3].axis('off') plt.tight_layout() plt.savefig(os.path.join(save_dir, f'sample_{idx}_{i}_iou_{iou:.3f}.png'), dpi=150, bbox_inches='tight') plt.close() print(f'平均IoU: {np.mean(iou_list):.4f} (+/- {np.std(iou_list):.4f})') print(f'最差IoU的样本: {np.min(iou_list):.4f}') print(f'最佳IoU的样本: {np.max(iou_list):.4f}')

通过可视化错误叠加图，你可以清晰看到模型在哪里犯错：是误将阴影识别为水体（假阳性），还是漏掉了细小河流（假阴性）？这为下一步优化指明了方向。

6.2 常见问题、原因与解决方案速查表

在训练和使用模型时，你可能会遇到以下典型问题：

6.3 模型部署与推理优化初步

当得到一个满意的模型后，你可能希望将其应用到新的卫星影像上。这里涉及模型转换和批量推理。

import glob from PIL import Image def inference_on_new_images(model, model_path, image_folder, output_folder, device): """ 对新图像文件夹进行批量推理 """ # 加载训练好的最佳模型 checkpoint = torch.load(model_path, map_location=device) model.load_state_dict(checkpoint['model_state_dict']) model.to(device) model.eval() os.makedirs(output_folder, exist_ok=True) image_paths = glob.glob(os.path.join(image_folder, '*.png')) + glob.glob(os.path.join(image_folder, '*.jpg')) transform = val_transform # 使用和验证集相同的预处理 for img_path in image_paths: image_name = os.path.basename(img_path) # 读取并预处理 image = cv2.imread(img_path) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) original_size = image.shape[:2] # (H, W) # 如果图像尺寸与训练时不同，可能需要调整（如填充或缩放） # 这里假设新图像尺寸与训练时一致，否则需要适配 input_tensor = transform(image=image)['image'].unsqueeze(0).to(device) with torch.no_grad(): output = model(input_tensor) prob_map = torch.softmax(output, dim=1)[0, 1].cpu().numpy() # 水体类别概率图 pred_mask = (prob_map > 0.5).astype(np.uint8) * 255 # 转为0-255的二值图 # 保存结果 pred_img = Image.fromarray(pred_mask) pred_img.save(os.path.join(output_folder, f'pred_{image_name}')) print(f'Processed: {image_name}') # 如果需要将模型转换为ONNX或TorchScript以便部署 def export_model(model, model_path, sample_input, export_name='water_seg_model.onnx'): checkpoint = torch.load(model_path, map_location='cpu') model.load_state_dict(checkpoint['model_state_dict']) model.eval() torch.onnx.export(model, sample_input, export_name, input_names=['input'], output_names=['output'], dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}) print(f"Model exported to {export_name}")

对于大规模影像推理，需要考虑将模型部署到GPU上并使用半精度（fp16）推理来提升速度。此外，如果处理的是大幅卫星影像，需要采用滑动窗口预测并拼接结果，同时处理好窗口重叠处的接缝问题。

围绕这个“开箱即用”的水体分割数据集，从数据理解、环境搭建、模型训练、调优技巧到问题排查和初步部署，整个流程走下来，你会发现它确实是一个极佳的研究起点。它让你避开了遥感AI中最脏最累的“数据准备”环节，直抵模型研发的核心。在实际操作中，最大的挑战往往来自于数据与任务的匹配度——这个数据集的水体类型、地域分布、季节变化是否覆盖了你的目标应用场景？如果不够，你可能还需要在此基础上进行微调或增量学习。我个人的体会是，用好一个现成数据集的关键，不在于盲目跑出多高的分数，而在于通过它深刻理解整个分割任务的 pipeline，掌握调优和解决问题的通用方法，这样当你面对自己独有的、更复杂的遥感数据时，才能游刃有余。

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