☰
Python+Unet遥感图像语义分割实战:从环境搭建到训练调参全流程
2026/10/12 20:32:06 网站建设 项目流程

简介:本资源为基于 Python 与 Unet 模型的图像语义分割完整项目包,面向希望入门或进阶计算机视觉的开发者,尤其适合作为毕业设计、课程设计、大作业或工程实训的参考方案。包内共 24 个文件,以 14 张 png 图像样本、4 个 py 脚本、1 个 h5 训练权重、1 个 pptx 说明文档及 xml、ovr 等辅助文件为主,压缩包约 478.98MB,涵盖数据生成、模型训练、预测推理与结果可视化等环节。项目围绕 Unet 编码器—解码器结构展开,提供从数据集构建到模型落地的完整链路,读者可据此理解语义分割的数据标注、网络搭建、损失计算与精度评估思路,并直接复用脚本进行二次开发。目前已有 204 人学习下载,适合需要快速搭建分割实验环境、对照代码查漏补缺的初学者与进阶学习者。

1. 从一张遥感图说起:python+Unet 图像语义分割到底在解决什么

手里有一批无人机航拍图或者遥感影像,想自动把建筑、道路、植被、水体一块块抠出来,逐像素打上类别标签——这就是图像语义分割要干的事。传统做法靠阈值、纹理特征加分类器,换个季节、换片区域就得重新调参,泛化能力差得让人头疼。Unet 这套编码器-解码器结构之所以在分割圈子里经久不衰,核心在于它的跳跃连接把浅层的高分辨率细节和深层的语义信息拼在一起,小目标边缘不容易丢,医学影像、遥感解译、工业缺陷检测都能用。python 生态里 PyTorch 或 TensorFlow 都能快速搭出 Unet,配合 numpy、OpenCV 做数据预处理,整条链路门槛并不高。这篇面向的是想自己动手跑通一套分割流程的工程师和新手:从环境装起,到数据组织、模型搭建、训练调参、推理可视化,再到踩坑排查,一步步给到能抄的代码和参数。遥感图像语义分割这类任务对显存和标注质量要求不低,读完你应该能判断自己的数据值不值得上 Unet,以及怎么少走弯路。

2. 环境与数据:把 python+Unet 的输入管道先搭对

2.1 python 环境与依赖安装的稳妥路径

很多人卡在第一步不是模型,是环境。python 安装教程网上一搜一大把,但真到跑 Unet 的时候,版本冲突才是血泪经验的来源。我一般用 conda 建独立环境,避免和系统里其他 python 项目打架。python 3.8 到 3.10 对主流深度学习框架兼容性最好,太新的版本有些库还没跟上。

# 创建独立环境,指定 python 3.9 conda create -n unet_seg python=3.9 -y conda activate unet_seg # 安装 PyTorch,按自己 CUDA 版本去官网查对应命令,这里以 CUDA 11.8 为例 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 # 图像处理与科学计算常用库 pip install numpy opencv-python pillow matplotlib scikit-learn tqdm

这段命令的逻辑是先隔离环境再装框架,最后补图像处理库。参数上注意--index-url要和你机器上的 CUDA 驱动匹配,装完用torch.cuda.is_available()验证,返回 False 说明 CUDA 没接上,别急着往下跑训练。numpy 是很多库的底层依赖,python 安装 numpy 库的方法虽然简单,但版本太新有时和 opencv 编译版本对不上,遇到ImportError先查 numpy 版本。vscode python 环境配置或者 pycharm 都行,关键是指定解释器到刚建的 conda 环境,不然跑起来用的还是系统 python。

2.2 语义分割数据集的组织与标注格式

Unet 训练需要成对的「原图 + 掩码图」。掩码图里每个像素的值就是类别索引,背景为 0,目标从 1 开始编号。遥感图像语义分割常见做法是把大图裁成 256×256 或 512×512 的小块,裁的时候要保证掩码和原图严格对齐,否则训练 loss 降不下去还找不到原因。

import os import cv2 import numpy as np def split_image_and_mask(img_path, mask_path, out_dir, patch_size=256, stride=128): """滑动窗口裁剪原图与掩码,stride 小于 patch_size 可做重叠增强""" img = cv2.imread(img_path) mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) h, w = img.shape[:2] idx = 0 for y in range(0, h - patch_size + 1, stride): for x in range(0, w - patch_size + 1, stride): img_patch = img[y:y+patch_size, x:x+patch_size] mask_patch = mask[y:y+patch_size, x:x+patch_size] # 跳过全背景的块,减少样本不均衡 if mask_patch.max() == 0: continue cv2.imwrite(os.path.join(out_dir, 'img', f'{idx}.png'), img_patch) cv2.imwrite(os.path.join(out_dir, 'mask', f'{idx}.png'), mask_patch) idx += 1 split_image_and_mask('raw/img.tif', 'raw/mask.png', 'dataset/train')

裁剪逻辑里patch_size决定单样本视野,太小语义信息不足,太大显存吃不消,256 是遥感任务的常用起点。stride设成 patch_size 的一半能增加样本量,但重叠太多会让训练集和验证集泄漏,划分时要以原图为单位切分,不能按 patch 随机分。掩码必须是单通道灰度图,类别值用整数,别用调色板 PNG,否则读出来是 RGB 三通道,训练直接报维度错误。

2.3 数据增强与 DataLoader 的落地写法

分割任务标注成本高,样本往往就几百张,数据增强是刚需。但增强必须原图和掩码同步做几何变换,翻转、旋转、缩放都要一致,颜色抖动只作用在原图上。

import torch from torch.utils.data import Dataset, DataLoader import albumentations as A from albumentations.pytorch import ToTensorV2 train_transform = A.Compose([ A.HorizontalFlip(p=0.5), A.VerticalFlip(p=0.5), A.RandomRotate90(p=0.5), A.RandomBrightnessContrast(p=0.3), A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), ToTensorV2(), ]) class SegDataset(Dataset): def __init__(self, img_dir, mask_dir, transform=None): self.img_dir = img_dir self.mask_dir = mask_dir self.transform = transform self.names = os.listdir(img_dir) def __len__(self): return len(self.names) def __getitem__(self, i): name = self.names[i] img = cv2.cvtColor(cv2.imread(os.path.join(self.img_dir, name)), cv2.COLOR_BGR2RGB) mask = cv2.imread(os.path.join(self.mask_dir, name), cv2.IMREAD_GRAYSCALE) if self.transform: aug = self.transform(image=img, mask=mask) img, mask = aug['image'], aug['mask'] return img, mask.long() loader = DataLoader(SegDataset('dataset/train/img', 'dataset/train/mask', train_transform), batch_size=8, shuffle=True, num_workers=4, pin_memory=True)

albumentations 的 Compose 保证几何变换同步作用于 image 和 mask,这是它比手写增强靠谱的地方。Normalize用的 ImageNet 均值方差,如果从头训练可以换成自己数据集的统计值。batch_size受显存限制,8 是 8G 显存的保守值,num_workers在 Windows 上设 0 更稳,Linux 可以开到 CPU 核数。mask 转long()是因为交叉熵损失要求标签是 int64。

3. Unet 网络结构:编码器、解码器与跳跃连接怎么写

3.1 从 Unet 网络结构图看每一层的职责

Unet 网络结构图那条经典的 U 形不是画着好看。左半边编码器每经过一次下采样,特征图尺寸减半、通道翻倍,感受野变大,负责回答「这是什么」。右半边解码器上采样恢复分辨率,负责回答「在哪里」。中间的跳跃连接把编码器同层特征直接拼到解码器对应层,让边缘、纹理这些空间细节不被深层卷积抹掉。遥感图像里道路又细又长,没有跳跃连接,分割结果经常断成一段一段。

import torch.nn as nn import torch.nn.functional as F 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), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, 3, padding=1), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), ) def forward(self, x): return self.block(x) class Down(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.pool_conv = nn.Sequential(nn.MaxPool2d(2), DoubleConv(in_ch, out_ch)) def forward(self, x): return self.pool_conv(x) class Up(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.up = nn.ConvTranspose2d(in_ch, in_ch // 2, 2, stride=2) self.conv = DoubleConv(in_ch, out_ch) def forward(self, x, skip): x = self.up(x) # 处理尺寸不整除导致的偏差 diff_y = skip.size(2) - x.size(2) diff_x = skip.size(3) - x.size(3) x = F.pad(x, [diff_x // 2, diff_x - diff_x // 2, diff_y // 2, diff_y - diff_y // 2]) x = torch.cat([skip, x], dim=1) return self.conv(x)

DoubleConv是 Unet 的基本砖块,两次 3×3 卷积加 BN 和 ReLU。Down用最大池化下采样,Up用转置卷积上采样后和跳跃特征拼接。拼接前那段 pad 是很多人忽略的坑:输入尺寸不是 16 的整数倍时,上采样后的尺寸和跳跃特征对不上,torch.cat直接报错,补零对齐是最省事的解法。

3.2 完整 Unet 组装与输出层通道数的确定

把砖块拼成完整网络,输出层通道数等于类别数(含背景)。遥感图像语义分割如果分建筑、道路、植被、水体加背景,就是 5 通道。

class UNet(nn.Module): def __init__(self, in_ch=3, n_classes=5): super().__init__() self.inc = DoubleConv(in_ch, 64) self.down1 = Down(64, 128) self.down2 = Down(128, 256) self.down3 = Down(256, 512) self.down4 = Down(512, 1024) self.up1 = Up(1024, 512) self.up2 = Up(512, 256) self.up3 = Up(256, 128) self.up4 = Up(128, 64) self.out = nn.Conv2d(64, n_classes, 1) def forward(self, x): x1 = self.inc(x) x2 = self.down1(x1) x3 = self.down2(x2) x4 = self.down3(x3) x5 = self.down4(x4) x = self.up1(x5, x4) x = self.up2(x, x3) x = self.up3(x, x2) x = self.up4(x, x1) return self.out(x)

n_classes一定要和掩码里的最大类别值加一对应,少一个通道训练时交叉熵会报 index 越界。in_ch按输入通道改,RGB 是 3,多光谱影像可能是 4 或更多。这个版本是最朴素的 Unet,unet 模型改进通常从这里下手:换编码器为 ResNet 或 EfficientNet 做迁移学习,加注意力模块,或者把转置卷积换成双线性插值上采样减少棋盘伪影。

3.3 损失函数与评价指标的选型理由

分割任务样本极不均衡,背景像素往往占九成以上,纯交叉熵会让模型学会全预测背景也能拿高准确率。常见做法是交叉熵加 Dice 损失组合。

class DiceLoss(nn.Module): def __init__(self, smooth=1e-6): super().__init__() self.smooth = smooth def forward(self, logits, targets): probs = F.softmax(logits, dim=1) targets_onehot = F.one_hot(targets, num_classes=logits.shape[1]).permute(0, 3, 1, 2).float() intersection = (probs * targets_onehot).sum(dim=(2, 3)) union = probs.sum(dim=(2, 3)) + targets_onehot.sum(dim=(2, 3)) dice = (2 * intersection + self.smooth) / (union + self.smooth) return 1 - dice.mean() # 组合损失 ce = nn.CrossEntropyLoss() def criterion(logits, targets): return ce(logits, targets) + DiceLoss()(logits, targets)

Dice 直接优化预测和真值的重叠度,对小目标更友好。smooth防止分母为零。评价指标别只看像素准确率,用 mIoU(平均交并比)更能反映真实分割质量,每个类别算 IoU 再平均,背景类通常不计入或单独看。

4. 训练、推理与可视化:把模型真正跑起来

4.1 训练循环与学习率、batch size 的调参经验

训练循环本身不复杂,难的是参数怎么设。我一般用 Adam 起步,学习率 1e-3 到 1e-4,配合余弦退火或 ReduceLROnPlateau。

import torch from tqdm import tqdm device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = UNet(in_ch=3, n_classes=5).to(device) optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', patience=5, factor=0.5) best_iou = 0 for epoch in range(100): model.train() loop = tqdm(loader, desc=f'Epoch {epoch}') for imgs, masks in loop: imgs, masks = imgs.to(device), masks.to(device) optimizer.zero_grad() logits = model(imgs) loss = criterion(logits, masks) loss.backward() optimizer.step() loop.set_postfix(loss=loss.item()) # 每个 epoch 后在验证集算 mIoU,这里省略验证函数细节 val_iou = evaluate(model, val_loader, device) scheduler.step(-val_iou) if val_iou > best_iou: best_iou = val_iou torch.save(model.state_dict(), 'best_unet.pth')

学习率太大 loss 震荡不降,太小收敛慢到怀疑人生,1e-3 是 Adam 的常用起点。patience=5表示验证指标 5 个 epoch 不提升就砍半学习率。batch size 和显存直接挂钩,显存不够优先降 batch 而不是降分辨率,分辨率对分割精度影响更大。训练时盯着验证集 mIoU,训练 loss 一直降但验证 mIoU 掉头向下就是过拟合,加增强或加 dropout。

4.2 推理脚本与分割结果可视化

训练完要能对新图出结果,推理时记得切 eval 模式和 no_grad。

def predict(model, img_path, device, n_classes=5): model.eval() img = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB) img = cv2.resize(img, (256, 256)) tensor = torch.from_numpy(img).permute(2, 0, 1).float().unsqueeze(0) / 255.0 tensor = tensor.to(device) with torch.no_grad(): logits = model(tensor) pred = torch.argmax(logits, dim=1).squeeze().cpu().numpy() return pred pred = predict(model, 'test.png', device) plt.imshow(pred, cmap='jet') plt.colorbar() plt.savefig('pred_vis.png', dpi=150, bbox_inches='tight')

argmax取每个像素概率最大的类别得到分割图。可视化时 python 画图横坐标太密集是常见问题,plt.xticks(rotation=45)或者直接关掉坐标轴plt.axis('off')更适合看分割效果。类别多的时候用固定颜色映射表,别用 jet,颜色跳变容易误读。

4.3 用 mIoU 和混淆矩阵验证分割质量

光看图不够,要量化。mIoU 按类别算,混淆矩阵能看出哪两类容易混。

def compute_iou(pred, target, n_classes): ious = [] for cls in range(n_classes): pred_mask = (pred == cls) target_mask = (target == cls) inter = (pred_mask & target_mask).sum() union = (pred_mask | target_mask).sum() if union == 0: continue ious.append(inter / union) return np.mean(ious)

union == 0的类别说明验证集里根本没出现,跳过不计入平均,否则会被拉低。混淆矩阵用 sklearn 的confusion_matrix展平后算,能直观看到建筑和道路是不是经常互相误判,这类错误往往要靠增加难例样本或调整损失权重来解决。

5. 避坑与排查:Unet 训练里最容易翻车的几件事

5.1 loss 不下降甚至变 NaN

现象:训练几个 batch 后 loss 变成 nan,或者一直卡在高位不动。原因通常是学习率过大、数据没归一化、或者掩码里有超出类别数的值。解决:先把学习率降到 1e-4 试,检查输入像素是否归一化到 0-1 或标准化,用np.unique(mask)确认掩码最大值小于 n_classes。混合精度训练时梯度溢出也会 NaN,加torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)做梯度裁剪。

5.2 验证集 mIoU 远低于训练集

现象:训练集 mIoU 0.9,验证集只有 0.4。原因多半是过拟合或数据泄漏。解决:先查数据划分,重叠裁剪的 patch 如果随机分到训练和验证就会泄漏,要以原图为单位划分。确认没泄漏后加数据增强、加权重衰减、或者换更轻的编码器。样本量太少时考虑迁移学习,用预训练编码器初始化。

5.3 分割结果边缘毛糙、小目标丢失

现象:大块区域分得还行,细长道路断裂,小建筑漏检。原因:下采样太多次丢失空间细节,或者损失函数对小目标不敏感。解决:减少下采样层数,或者用空洞卷积替代部分池化;损失里加大 Dice 权重;推理时用滑动窗口重叠预测再拼接,边缘处取重叠区域的平均概率。

5.4 显存溢出 OOM

现象:训练中途报 CUDA out of memory。原因:batch size 太大、分辨率太高、或者验证时没加 no_grad 导致计算图累积。解决:降 batch size,用梯度累积模拟大 batch;推理和验证包在torch.no_grad()里;及时del中间变量并torch.cuda.empty_cache()。混合精度训练也能省不少显存。

5.5 推理速度慢到无法落地

现象:单张图推理要好几秒。原因:模型太大、没转半精度、或者用了 CPU。解决:推理时model.half()转 FP16,输入 tensor 也转 half;导出 ONNX 或 TensorRT 加速;确认torch.cuda.is_available()为 True,别在 CPU 上硬跑。遥感图像语义分割如果要做大图推理,分块并行比整图输入更实际。

6. 进阶技巧:让 Unet 在你自己数据上真正好用

跑通基础版只是起点,真正决定效果的是几个细节。第一,编码器换预训练权重。从头训练 Unet 在小数据集上很难收敛好,把编码器换成在 ImageNet 上预训练的 ResNet34 或 EfficientNet-B0,解码器随机初始化,mIoU 通常能涨 5 到 10 个点。第二,损失函数按类别加权。遥感任务里道路、水体样本少,给这些类更高的交叉熵权重,或者用 Focal Loss 压制易分类样本。第三,测试时增强(TTA)。推理时对同一张图做水平翻转、垂直翻转各预测一次,概率图平均后再 argmax,边缘稳定性明显提升,代价是推理时间翻几倍,看场景取舍。

验证方法上,别只信一个 mIoU 数字。我会把验证集里 IoU 最低的几张图单独拎出来看,往往能发现标注错误或者某类目标系统性漏检。标注质量是分割任务的天花板,模型再强也救不了错标。另外,训练日志里除了 loss 和 mIoU,把每个类别的 IoU 都打出来,某一类突然掉下去通常对应数据里那类样本分布变了。

一个具体技巧:学习率预热。前 500 个 iteration 从 1e-6 线性升到设定值,再正常衰减,能明显减少训练初期的震荡,尤其是 batch size 较大的时候。代码上用一个 lambda 函数控制 optimizer 的 lr 就行。

def warmup_lr(step, warmup_steps=500, base_lr=1e-3): if step < warmup_steps: return base_lr * step / warmup_steps return base_lr for step, (imgs, masks) in enumerate(loader): lr = warmup_lr(step) for pg in optimizer.param_groups: pg['lr'] = lr # 正常训练步骤...

我自己踩得最狠的一次坑是掩码用了调色板 PNG,读出来是 RGB,训练 loss 死活不降,查了两天才发现是数据格式问题。从那以后我养成了一个习惯:拿到任何分割数据集,先np.unique看掩码值,再cv2.imread确认通道数,这两步花不了一分钟,能省掉大量玄学调试。希望帮到你。

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

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询