GeoSeg 框架解析

总的来看：

config 配置文件
geoseg 定义了核心组件：models datasets losses
tools 辅助工具箱

config :

里面包含多个遥感数据集，内部包含不同的模型

看一个例子

复制代码

导入
from torch.utils.data import DataLoader
from geoseg.losses import *
from geoseg.datasets.vaihingen_dataset import *
from geoseg.models.UNetFormer import UNetFormer
from tools.utils import Lookahead, process_model_params

训练超参数
from torch.utils.data import DataLoader
from geoseg.losses import *
from geoseg.datasets.vaihingen_dataset import *
from geoseg.models.UNetFormer import UNetFormer
from tools.utils import Lookahead, process_model_params

网络定义
net = UNetFormer(num_classes=num_classes)

损失函数
loss = UnetFormerLoss(ignore_index=ignore_index)
use_aux_loss = True

数据加载器
train_dataset = VaihingenDataset(data_root='data/vaihingen/train', ...)
val_dataset = VaihingenDataset(...)
...
train_loader = DataLoader(dataset=train_dataset, ...)
val_loader = DataLoader(dataset=val_dataset, ...)

优化器
layerwise_params = {"backbone.*": dict(lr=backbone_lr, ...)}
net_params = process_model_params(net, ...)
base_optimizer = torch.optim.AdamW(net_params, lr=lr, ...)
optimizer = Lookahead(base_optimizer)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(...)

geoseg

models--->所有神经网络骨架库

datasets-->数据集库

losses --->损失函数库

datasets

除了每个数据集的datasets 之外，有个transform是数据增强策略

transform: 各种各样的数据增强策略

看一个dataset例子：

复制代码

全局常量
类别
CLASSES = ('ImSurf', 'Building', 'LowVeg', 'Tree', 'Car', 'Clutter')
可视化颜色
PALETTE = [[255, 255, 255], [0, 0, 255], [0, 255, 255], [0, 255, 0], [255, 204, 0], [255, 0, 0]]

ORIGIN_IMG_SIZE = (1024, 1024)
INPUT_IMG_SIZE = (1024, 1024)
TEST_IMG_SIZE = (1024, 1024)

数据增强流水线
def get_training_transform(): ...
def train_aug(img, mask): ...
def get_val_transform(): ...
def val_aug(img, mask): ...
这些都是对数据进行数据增强

核心数据集类
class VaihingenDataset(Dataset):
    def __init__(...): ...构造函数，保存传入参数，调用get_img_ids读取images_1024目录下的所有文件名，并且保存在self.img_ids 当中
    def __getitem__(...): ...
    def __len__(...): ...
    def get_img_ids(...): ...用来确保img_dir 和mask_dir能对应的上
    def load_img_and_mask(...): ...加载数据
    def load_mosaic_img_and_mask(...): ...高级数据增强
    
可视化工具
def show_img_mask_seg(...): ...
def show_seg(...): ...
def show_mask(...): ...

losses

各种各样loss 函数

models：定义了不同的models 神经网络

看一个例子

这个例子没有encoder 模块是因为，他的encoder 模块就只是backbone 特征提取，如果特征提取后加入混合注意力之类的东西，就应该写一个encoder类

此外，models 一般都是把组件，基础架构模块定义好，核心模块定义好，在encoder 和 decoder 内部，按照论文结构排列组合
encoder 和 decoder 会在最后的模型封装中，在forward 中应用

复制代码

基础架构模块
class ConvBNReLU(nn.Sequential): ...
class ConvBN(nn.Sequential): ...
class Conv(nn.Sequential): ...
class SeparableConvBNReLU(nn.Sequential): ...
class Mlp(nn.Module): ...

核心注意力模块
class GlobalLocalAttention(nn.Module): ...

Transformer 模块
class GlobalLocalAttention(nn.Module): ...

解码辅助模块
class WF(nn.Module): ...
class FeatureRefinementHead(nn.Module): ...
class AuxHead(nn.Module): ...

解码器
class Decoder(nn.Module): ...

最终模型
class UNetFormer(nn.Module): ...
构造函数和前向传播

tools

数据预处理脚本

模型配置脚本

模型评估脚本