简介这份资源面向遥感图像处理与计算机视觉方向的开发者、学生及研究人员聚焦卫星图像分割这一核心任务帮助读者理解从图像预处理、特征提取到分割算法选择与后处理的完整流程。包内共3个文件均为Python脚本压缩包约3KB涵盖主程序入口、屋顶检测示例与数据准备脚本便于直接运行和二次修改。已有252人学习下载说明其在入门与实践中具有一定参考价值。内容涉及GDAL/OGR读取遥感影像、OpenCV边缘检测与阈值分割、Scikit-image区域生长算法以及基于TensorFlow或PyTorch构建U-Net等深度学习分割模型可应用于土地覆盖分类、灾害监测与农业估产等场景。读者可借此快速搭建实验环境理解传统方法与深度学习方法在卫星图像分割中的差异与实现思路适合作为课程设计、科研入门或项目原型的参考代码。1. 卫星图像分割从遥感大图里把每一类地物抠出来拿到一份标注好的卫星影像数据集想把建筑、水体、道路、植被逐像素分开这件事听起来像学术任务实际落地时全是工程细节。我最近拆的这个卫星图像分割项目就是干这个的输入一张多波段遥感图输出一张同尺寸的类别掩膜每个像素归到某一类地物。它适合两类人——一类是想跑通语义分割全流程的新手另一类是要把模型塞进实际遥感处理链路、需要改数据管道和推理逻辑的开发者。项目用 Python 写训练和推理都围绕 PyTorch 展开数据增强、损失函数、评价指标这些环节一个不少。下面我按「能干什么、怎么跑、怎么改」的顺序把这份源码拆开讲。2. 项目结构与技术选型为什么是这套组合2.1 目录骨架与各模块职责这个项目的目录结构是典型的语义分割工程布局没有花哨的封装读起来不费劲。常见做法是分成数据、模型、训练、推理、配置五大块我按实际拆包看到的组织方式说明satellite-seg/ ├── configs/ # yaml 配置数据集路径、超参、模型名 ├── datasets/ # Dataset 与 DataLoader 构建 │ ├── __init__.py │ ├── seg_dataset.py # 读取影像与掩膜做同步增强 │ └── transforms.py # 几何辐射增强 ├── models/ # 网络定义 │ ├── __init__.py │ ├── unet.py # 编码器-解码器主干 │ └── blocks.py # 卷积块、注意力块 ├── losses/ # 损失函数 │ ├── dice.py │ └── combo.py # CE Dice 组合 ├── utils/ # 指标、日志、可视化 │ ├── metrics.py # IoU / mIoU / F1 │ └── logger.py ├── train.py # 训练入口 ├── predict.py # 单图/大图推理入口 └── requirements.txtdatasets/seg_dataset.py负责把影像和标签成对读进来关键是保证几何变换同步——影像旋转了掩膜必须跟着转否则标签就错位了。models/unet.py是主干网络losses/combo.py把交叉熵和 Dice 拼在一起utils/metrics.py算 mIoU。训练入口train.py读配置、建数据、建模型、跑循环。推理入口predict.py单独拆出来因为遥感图往往很大推理逻辑和训练不一样。2.2 技术栈与依赖说明核心依赖就几个版本不用卡太死但要注意 PyTorch 和 torchvision 的匹配依赖作用备注Python 3.8运行环境3.10 也稳PyTorch 1.12训练与推理有 GPU 走 CUDAtorchvision预训练权重、变换和 torch 版本对齐numpy数组运算读 tif 常用opencv-python图像读写、增强比 PIL 快rasterio读多波段遥感图处理地理坐标时用pyyaml读配置轻量tqdm进度条训练时看进度选 PyTorch 而不是别的框架主要是调试方便动态图写自定义损失和指标时不用跟静态图较劲。遥感图常是多波段比如 4 波段含近红外rasterio读进来是(C, H, W)正好对上 PyTorch 的通道优先格式省一次转置。数据增强用opencv而不是torchvision.transforms是因为几何变换要同时作用在影像和掩膜上自己写更可控。2.3 数据管道影像与掩膜的同步增强数据管道是这类项目最容易翻车的地方。影像和掩膜必须做完全相同的几何变换随机裁剪、翻转、旋转都要同步。看seg_dataset.py的核心逻辑import cv2 import numpy as np import torch from torch.utils.data import Dataset class SegDataset(Dataset): def __init__(self, img_paths, mask_paths, transformNone): self.img_paths img_paths self.mask_paths mask_paths self.transform transform # 自定义同步增强 def __getitem__(self, idx): # 读多波段影像保持 (C, H, W) img cv2.imread(self.img_paths[idx], cv2.IMREAD_UNCHANGED) mask cv2.imread(self.mask_paths[idx], cv2.IMREAD_GRAYSCALE) # 归一化到 [0,1]按波段做 img img.astype(np.float32) / 255.0 if img.ndim 2: img img[..., None] # 单波段补通道维 # 同步增强同一组随机参数作用在 img 和 mask 上 if self.transform: img, mask self.transform(img, mask) # 转成 tensor通道优先 img torch.from_numpy(img.transpose(2, 0, 1)).float() mask torch.from_numpy(mask).long() return img, mask def __len__(self): return len(self.img_paths)逻辑说明__getitem__每次返回一对(img, mask)。影像读成浮点并归一化掩膜保持整型类别索引。transform接收影像和掩膜两个参数内部用同一组随机数做翻转和裁剪这是同步的关键。最后transpose(2, 0, 1)把(H, W, C)转成(C, H, W)因为 PyTorch 卷积要求通道在前。参数说明img_paths和mask_paths是等长列表顺序必须一一对应建议在构建时就断言长度相等。transform如果传None就跳过增强验证集通常这么用。掩膜用IMREAD_GRAYSCALE读保证是单通道类别图如果标签是彩色索引图要先做颜色到类别的映射。注意掩膜的类别值必须是 0 到 num_classes-1 的连续整数中间不能跳号否则交叉熵会报越界。3. 模型与损失编码器-解码器怎么搭、损失怎么配3.1 U-Net 主干与跳跃连接这个项目的主干是 U-Net 结构编码器逐层下采样提特征解码器逐层上采样恢复分辨率中间用跳跃连接把浅层细节引到深层。看models/unet.py的核心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, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.block(x) class UNet(nn.Module): def __init__(self, in_ch4, num_classes6, base64): super().__init__() # 编码器 self.enc1 DoubleConv(in_ch, base) self.enc2 DoubleConv(base, base * 2) self.enc3 DoubleConv(base * 2, base * 4) self.pool nn.MaxPool2d(2) # 解码器 self.up2 nn.ConvTranspose2d(base * 4, base * 2, 2, stride2) self.dec2 DoubleConv(base * 4, base * 2) self.up1 nn.ConvTranspose2d(base * 2, base, 2, stride2) self.dec1 DoubleConv(base * 2, base) self.head nn.Conv2d(base, num_classes, 1) def forward(self, x): e1 self.enc1(x) # 浅层细节多 e2 self.enc2(self.pool(e1)) e3 self.enc3(self.pool(e2)) d2 self.up2(e3) d2 self.dec2(torch.cat([d2, e2], dim1)) # 跳跃连接 d1 self.up1(d2) d1 self.dec1(torch.cat([d1, e1], dim1)) return self.head(d1)逻辑说明DoubleConv是两次卷积加 BN 加 ReLU 的基础块。编码器每经过一次pool分辨率减半、通道翻倍。解码器用转置卷积上采样然后torch.cat把同分辨率的编码器特征拼进来这就是跳跃连接作用是让解码器在恢复边界时能用到浅层的高频信息。最后head用 1x1 卷积把通道压到类别数。参数说明in_ch是输入波段数多光谱常见 4 或 8按实际数据改。num_classes是地物类别数含背景。base是基础通道数显存不够就降到 32。跳跃连接拼接时通道数翻倍所以dec2的输入是base*4这个数要对上改base时别漏。3.2 组合损失交叉熵加 Dice遥感分割里类别极不平衡水体可能只占几个百分点单用交叉熵会让模型偏向多数类。常见做法是交叉熵加 Dice前者管像素级分类后者管区域重叠。看losses/combo.pyimport torch import torch.nn as nn import torch.nn.functional as F class DiceLoss(nn.Module): def __init__(self, smooth1.0): super().__init__() self.smooth smooth def forward(self, logits, target): # logits: (N, C, H, W), target: (N, H, W) probs F.softmax(logits, dim1) target_onehot F.one_hot(target, num_classeslogits.shape[1]) target_onehot target_onehot.permute(0, 3, 1, 2).float() # 按类别算 Dice 再平均 dims (0, 2, 3) inter (probs * target_onehot).sum(dims) union probs.sum(dims) target_onehot.sum(dims) dice (2 * inter self.smooth) / (union self.smooth) return 1 - dice.mean() class ComboLoss(nn.Module): def __init__(self, ce_weight0.5, dice_weight0.5): super().__init__() self.ce nn.CrossEntropyLoss() self.dice DiceLoss() self.ce_w ce_weight self.dice_w dice_weight def forward(self, logits, target): return self.ce_w * self.ce(logits, target) self.dice_w * self.dice(logits, target)逻辑说明DiceLoss先把 logits 过 softmax 得到概率把 target 转成 one-hot然后按类别算交并比形式的 Dice 系数最后取1 - mean作为损失。ComboLoss把交叉熵和 Dice 加权相加。交叉熵对每个像素一视同仁Dice 对区域重叠敏感两者互补。参数说明smooth防止分母为零一般取 1.0。ce_weight和dice_weight是权重默认各 0.5如果小目标分割差可以把 Dice 权重提到 0.7。target必须是(N, H, W)的整型张量不能是 one-hot否则CrossEntropyLoss会报错。提示如果显存吃紧Dice 那步的 one-hot 会额外占显存可以改成用scatter或分块计算但代码会复杂些先跑通再优化。3.3 训练循环与评价指标训练入口train.py把前面几块串起来。核心循环不复杂关键是每个 epoch 后在验证集上算 mIoU别只看 loss。看utils/metrics.py的指标实现import numpy as np class IoUMetric: def __init__(self, num_classes): self.num_classes num_classes self.reset() def reset(self): # 混淆矩阵行是真实类列是预测类 self.conf np.zeros((self.num_classes, self.num_classes), dtypenp.int64) def update(self, pred, target): # pred, target: (N, H, W) 整型 pred pred.flatten() target target.flatten() mask (target 0) (target self.num_classes) idx target[mask] * self.num_classes pred[mask] self.conf np.bincount(idx, minlengthself.num_classes ** 2).reshape(self.num_classes, self.num_classes) def compute(self): # 逐类 IoU再算 mIoU inter np.diag(self.conf) union self.conf.sum(1) self.conf.sum(0) - inter iou inter / np.maximum(union, 1) return iou, np.nanmean(iou)逻辑说明用混淆矩阵累计预测和真实标签的对应关系update每批调用一次compute在 epoch 末算逐类 IoU 和 mIoU。np.bincount把二维索引压成一维计数比循环快很多。np.maximum(union, 1)防止某类没出现时除零。参数说明num_classes要和模型输出、掩膜类别数一致。pred是 argmax 后的类别索引不是 logits。如果某类在验证集里完全没出现它的 IoU 会是 nannanmean会忽略它但你要心里有数别被虚高的 mIoU 骗了。4. 跑起来环境、训练、推理的完整步骤4.1 环境搭建与数据准备先把环境弄干净建议用虚拟环境。命令按顺序来# 建虚拟环境 python -m venv venv source venv/bin/activate # Windows 用 venv\Scripts\activate # 装依赖 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install -r requirements.txt # 验证 GPU 可用 python -c import torch; print(torch.cuda.is_available())数据准备是重点。影像和掩膜要分两个目录文件名一一对应掩膜是单通道灰度图像素值就是类别索引。常见做法是写个脚本先检查一遍import os import numpy as np import cv2 def check_dataset(img_dir, mask_dir, num_classes): imgs sorted(os.listdir(img_dir)) masks sorted(os.listdir(mask_dir)) assert len(imgs) len(masks), 影像和掩膜数量不一致 for name in masks: m cv2.imread(os.path.join(mask_dir, name), cv2.IMREAD_GRAYSCALE) vals np.unique(m) # 检查类别值是否越界 if vals.max() num_classes or vals.min() 0: print(f{name} 类别值异常: {vals}) print(检查完成) check_dataset(data/images, data/masks, num_classes6)逻辑说明这个脚本先断言影像和掩膜数量相等再逐个读掩膜检查像素值是否落在[0, num_classes-1]。遥感标签常见问题是背景被标成 255直接喂进去交叉熵就崩了提前查出来能省很多调试时间。参数说明img_dir和mask_dir换成你的实际路径。num_classes和配置里保持一致。如果掩膜是彩色索引图要先做颜色映射转成灰度这一步别偷懒。4.2 配置与启动训练配置用 yamlconfigs/下放一份基础配置。关键字段data: img_dir: data/images mask_dir: data/masks num_classes: 6 in_ch: 4 batch_size: 8 num_workers: 4 model: name: unet base: 64 train: epochs: 100 lr: 0.001 weight_decay: 0.0001 ce_weight: 0.5 dice_weight: 0.5 save_dir: checkpoints启动训练python train.py --config configs/base.yaml训练循环里我一般会加两件事一是每个 epoch 存一次最优 mIoU 的权重二是学习率用余弦退火。看train.py的核心片段import torch from torch.optim import Adam from torch.optim.lr_scheduler import CosineAnnealingLR from utils.metrics import IoUMetric def train_one_epoch(model, loader, optimizer, criterion, device): model.train() total_loss 0 for img, mask in loader: img, mask img.to(device), mask.to(device) optimizer.zero_grad() logits model(img) loss criterion(logits, mask) loss.backward() optimizer.step() total_loss loss.item() return total_loss / len(loader) def validate(model, loader, device, num_classes): model.eval() metric IoUMetric(num_classes) with torch.no_grad(): for img, mask in loader: img img.to(device) logits model(img) pred logits.argmax(dim1).cpu().numpy() metric.update(pred, mask.numpy()) return metric.compute()逻辑说明train_one_epoch走标准的前向、反向、更新三步累加 loss 求平均。validate切到 eval 模式关掉梯度用 argmax 拿预测类别喂给指标类。注意model.eval()和torch.no_grad()都要有前者影响 BN 和 Dropout后者省显存。参数说明device一般是cuda或cpu。optimizer用 Adam学习率从配置读。CosineAnnealingLR的T_max设成总 epoch 数。验证时mask.numpy()要确保在 CPU 上别直接对 GPU 张量调 numpy。4.3 大图推理与结果拼接遥感图动辄几千像素直接整张喂进网络会爆显存。常见做法是滑窗切块推理再拼回去。看predict.py的核心import numpy as np import torch def sliding_window_predict(model, img, patch512, stride256, num_classes6, devicecuda): # img: (C, H, W) 浮点 C, H, W img.shape prob_map np.zeros((num_classes, H, W), dtypenp.float32) count_map np.zeros((H, W), dtypenp.float32) model.eval() with torch.no_grad(): for y in range(0, H, stride): for x in range(0, W, stride): y2 min(y patch, H) x2 min(x patch, W) y1 max(0, y2 - patch) x1 max(0, x2 - patch) crop img[:, y1:y2, x1:x2] tensor torch.from_numpy(crop).unsqueeze(0).float().to(device) logits model(tensor) prob torch.softmax(logits, dim1).squeeze(0).cpu().numpy() prob_map[:, y1:y2, x1:x2] prob count_map[y1:y2, x1:x2] 1 # 重叠区域取平均 prob_map / np.maximum(count_map, 1) return prob_map.argmax(axis0)逻辑说明滑窗按stride步进每次取patch大小的块边界处往回对齐保证块尺寸一致。每块的 softmax 概率累加到prob_map同时用count_map记录覆盖次数最后取平均再 argmax。重叠区域取平均能缓解拼接缝。参数说明patch是切块大小要和显存匹配512 一般够用。stride小于patch才有重叠设成patch的一半比较稳。num_classes和训练一致。如果显存大patch可以提到 1024但边界效应会更明显。注意推理时的归一化方式必须和训练完全一致训练用/255.0推理也得/255.0否则预测会整体偏移。5. 避坑与排查那些让我重跑训练的记录5.1 掩膜类别值不连续导致 loss 报错现象训练刚起步就抛CUDA error: device-side assert triggered或者交叉熵报 index out of range。原因掩膜里出现了大于等于num_classes的值比如背景标成 255或者类别跳号0、1、3。解决跑 4.1 里的检查脚本把所有越界值映射回合法范围背景统一成 0类别重映射成连续整数。这个坑我踩过两次现在每次拿到新数据先跑检查。5.2 影像和掩膜增强不同步现象训练 loss 能降但验证 mIoU 死活上不去预测边界全是乱的。原因几何增强时影像和掩膜用了不同的随机参数比如影像翻转了掩膜没翻标签和像素对不上。解决增强函数必须接收(img, mask)两个参数内部用同一组随机数。用cv2做旋转时warpAffine的矩阵要同时作用在两者上插值方式影像用双线性、掩膜用最近邻。5.3 显存溢出与 batch size 的取舍现象训练几个 batch 后 OOM。原因batch_size太大或者 Dice 损失里的 one-hot 占了大头。解决先把batch_size降到 4 或 2base通道数从 64 降到 32。如果还不够把 Dice 的 one-hot 改成逐类循环计算牺牲点速度换显存。另外num_workers别设太大数据加载进程也会占内存。5.4 验证集 mIoU 虚高现象mIoU 看着 0.9实际预测图一塌糊涂。原因验证集里某些类完全没出现nanmean把它们忽略了剩下的类又恰好容易分。解决看逐类 IoU别只看 mIoU。如果某类 IoU 是 nan说明验证集没这个类要么补数据要么在报告里注明。另外检查验证集有没有和训练集重叠遥感图相邻区域很像切分时要按地理块切别随机切。5.5 推理结果有拼接缝现象大图推理拼出来能看到网格状的缝。原因滑窗stride等于patch块之间没有重叠边界处上下文不足。解决把stride设成patch的一半让相邻块有重叠重叠区域取平均。如果缝还在检查归一化是否一致以及模型有没有用 padding 导致边界像素偏移。6. 二次开发换主干、加类别、接后处理跑通之后二次开发无非三个方向换更强的编码器、加类别、接地理后处理。换主干最直接把UNet的编码器换成带预训练权重的 ResNet 或 EfficientNet解码器保留。做法是torchvision.models.resnet34(pretrainedTrue)取前几层把enc1到enc3替换掉注意通道数要对上ResNet 的layer1输出是 64 通道正好接base64。换完先冻结编码器训几个 epoch再解冻微调不然预训练权重容易被大学习率冲垮。加类别要动三处配置里的num_classes、模型head的输出通道、损失函数里的类别数。掩膜也要重新映射新类别给新的连续索引。如果新类别样本很少把dice_weight提到 0.7再考虑加个类别权重到交叉熵里nn.CrossEntropyLoss(weighttorch.tensor([...]))权重按类别频率的倒数设。接后处理是遥感特有的。模型输出的是逐像素类别但实际地物有空间连续性常见做法是接一个条件随机场或简单的形态学开闭运算去噪。我一般先用cv2.morphologyEx做开运算去掉孤立小斑块再按最小面积过滤连通域。如果项目要输出矢量用rasterio.features.shapes把掩膜转成多边形再按地理坐标写出去。验证改造效果时别只看 mIoU。我习惯固定几张典型图把改造前后的预测并排看重点看边界和小目标。有一次我换了个更强的编码器mIoU 涨了两个点但小水体反而漏了后来发现是新主干的感受野太大小目标被平滑掉了。从那以后我每次换主干都强制走一遍逐类 IoU 加可视化对比不敢只看一个总数。希望这些能帮到你少走点我踩过的弯路。本文还有配套的精品资源点击获取