MMSegmentation入门

1,240 阅读19分钟

MMSegmentation 是 OpenMMLab 开源的基于 PyTorch 实现的功能强大的语义分割工具箱 \

MMSegmentation 的主要特性如下:

  • 丰富的语义分割模型: 已支持 11 种主干网络和 34 种算法,例如常用模型 FCN, PSPNet 和 DeepLabV3;Transformer 模型,Swin Transformer、Segmenter 和 SegFormer; Real-Time 实时分割模型, ICNet、BiSeNet 和 STDC 等;以及最近流行的网络 ConvNeXt 和 MAE。(具体支持哪些网络,可查看官方文档)
  • 大量开箱即用的模型权重:在 16 个常用的语义分割数据集上提供了 590 个训练好的模型。
  • 统一的性能评估框架:优化和统一了训练和测试的流程,方便公平比较各个模型在特定任务上的表现。

目前,MMSegmentation 支持的分割任务为语义分割 ,MMDetection 中支持了实例分割和全景分割。

MMSegmentation 算法库框架介绍

本文基于MMSegmentation v0.23.0 的整体架构(2022.05.26)
具体当前已实现的框架得结合当前最新代码及文档!

MMSegmentation 目录结构

按照代码目录下的文件夹,MMSegmentation 代码库主要可以包含四个部分: 1)./tools 包括了调用 MMSegmentation 作为训练和测试入口的./tools/train.py 和 ./tools/test.py,预训练模型和数据集准备的转换脚本,以及部署和可视化相关的脚本。 详细介绍可见 Github 里的文档 2)./configs 包括了各个算法的配置文件、存放常用的数据集配置、基础模型以及训练策略的基配置文件 ./configs/_base_。 3)./mmseg 里面是 MMSegmentation 的算法库,包括核心组件、数据集处理、分割模型代码和面向用户的 API 接口。 4)./data 指的是存放数据集的路径,在原本的代码库中没有这个文件夹。用户只需指定正确的文件夹路径即可使用数据。
下面是详细的 MMSegmentation 的算法库目录结构:(具体可直接看代码)

# MMSegmentation 算法库目录结构的主要部分
mmsegmentation
    |
    |- configs                  #配置文件
    |- _base_                    ## 基配置文件
    |     |     |- datasets       ### 数据集相关配置文件
    |     |     |- models               ### 模型相关配置文件
    |     |     |- schedules            ### 训练日程如优化器,学习率等相关配置文件
    |     |     |- default_runtime.py   ### 运行相关的默认的设置
    |     |- swin                     ## 各个分割模型的配置文件,会引用 _base_ 的配置并做修改
    |     |- ...                           
    |- data                           # 原始及转换后的数据集文件
    |- mmseg
    |     |- core                     ## 核心组件
    |     |     |- evaluation           ### 评估模型性能代码
    |     |- datasets                 ## 数据集相关代码
    |     |     |- pipelines            ### 数据预处理   
    |     |     |- samplers             ### 数据集采样代码   
    |     |     |- ade.py               ### 各个数据集准备需要的代码   
    |     |     |- ...   
    |     |- models                    ## 分割模型具体实现代码   
    |     |     |- backbones             ### 主干网络   
    |     |     |- decode_heads          ### 解码头   
    |     |     |- losses                ### 损失函数   
    |     |     |- necks                 ### 颈  
    |     |     |- segmentors            ### 构建完整分割网络的代码   
    |     |     |- utils                 ### 构建模型时的辅助工具   
    |     |- apis                      ## high level 用户接口,在这里调用 ./mmseg/ 内各个组件   
    |     |     |- train.py              ### 训练接口   
    |     |     |- test.py               ### 测试接口   
    |     |     |- ...   
    |     |- ops                       ## cuda 算子(即将迁移到 mmcv 中)   
    |     |- utils                     ## 辅助工具   
    |- tools   
    |     |- model_converters          ## 各个主干网络预训练模型转 key 脚本
    |     |- convert_datasets          ## 各个数据集准备转换脚本   
    |     |- train.py                  ## 训练脚本   
    |     |- test.py                   ## 测试脚本   
    |     |- ...                         
    |- ...

MMSegmentation 的算法库有 3 个关键组件:

1)./mmseg/apis/,用于训练和测试的接口。

2)./mmseg/models/,用于分割网络模型的具体实现。

3)./mmseg/datasets/,用于数据集处理。

本文我们主要介绍网络模型结构,因此涉及内容主要在 ./mmseg/models 里面。

MMSegmentation 模型实现

Segmentor

MMSegmentation 中将语义分割模型定义为 segmentor, 一般包括 backbone、neck、head、loss 4 个核心组件,每个模块的功能如下:

  1. 预处理后的数据输入到 backbone(如 ResNet 和 Swin Transformer)中进行编码并提取特征。
  2. 输出的单尺度或者多尺度特征图输入到 neck 模块中进行特征融合或者增强,典型的 neck 是 特征金字塔(Feature Pyramid Networks, FPN)。
  3. 上述多尺度特征最终输入到 head 部分,一般包括 decoder head, auxiliary head 以及 cascade decoder head,用以预测分割结果(它们的区别我们会在下文具体介绍)。
  4. 最后一步是计算 pixel 分类的 loss,进行训练。

Segmentor.png

需要说明的是,上述 4 个组件不是每个算法都需要的,比如很多模型里没有 neck 和 auxiliary head 组件。分割器 segmentor 的具体代码见文件 ./mmseg/models/segmentors/

MMSegmentation 里面的分割器框架可以分为 “Encoder Decoder” 结构和 “Cascade Encoder Decoder” 结构。现有的大多数模型为 “Encoder Decoder” 结构,即利用 encoder 提取图像特征,再用 decoder 去解码上述特征。 “Cascade Encoder Decoder” 的解码部分不是单独的解码头,而是级联式的 2 个或多个解码头,前一个解码头的输出作为后一个解码头的输入

关于分割器 segmentor 的训练和测试的基本逻辑,以语义分割经典的 “Encoder Decoder” 结构为例:

class EncoderDecoder(BaseSegmentor):
    def __init__(...):        
        # 构建 backbone、neck 和 head        
        self.backbone = build_backbone(backbone)        
        if neck is not None:            
            self.neck = build_neck(neck)
        self._init_decode_head(decode_head)
        self._init_auxiliary_head(auxiliary_head)  
    def forward_train(...):         
        # 利用 backbone+neck 进行特征提取        
        x = self.extract_feat(img)        
        losses = dict()        
        # decode head 输出预测特征图并计算出 loss        
        loss_decode = self._decode_head_forward_train(x, img_metas,                                                      gt_semantic_seg)
        losses.update(loss_decode)        
        # auxiliary heads 输出预测特征图并计算出 loss        
        if self.with_auxiliary_head:            
            loss_aux = self._auxiliary_head_forward_train(x, img_metas, gt_semantic_seg)            
            losses.update(loss_aux)        
        return losses          
    def simple_test(...):        
        # 调用 inference 函数,对输入图片做全图或者滑动窗口的推理,得到 logits 值        
        seg_logit = self.inference(img, img_meta, rescale)        
        # 做 argmax 得到预测的 prediction mask        
        seg_pred = seg_logit.argmax(dim=1)
    def aug_test(...):        
        ...

接下来,我们详细介绍分割器 segmentor 里 4 个核心组件:backbone, neck,head,和 loss。

Backbone

目前 MMSegmengtation 中已经集成了大部分主干网络,具体见文件 ./mmseg/models/backbones/,v0.24.1 已经实现的骨架如下:

a30cc6b5fb16d5edb686c206071bd264.png 通常定义的“主干网络” 是指从上游任务(如 ImageNet )预训练,然后用于多个下游任务(如目标检测、实例分割、语义分割、姿态估计)中的网络,而在 ./mmseg/models/backbones 里主干网络的定义有所不同,会把一些分割算法的网络结构也作为“主干网络”,如 UNet、 FastSCNN、CGNet、ICNet、BiSeNetV1/V2、ERFNet、STDC。

其中最常用的是 ResNet v1c 系列和 Vision Transformer 系列。如果你需要对骨架进行扩展,可以继承上述网络,然后通过注册器机制注册使用。一个典型用法为 ./configs/_base_/models/segmenter_vit-b16_mask.py 里面的:

checkpoint = 'https://download.openmmlab.com/mmsegmentation/v0.5/pretrain/segmenter/vit_base_p16_384_20220308-96dfe169.pth'  # noqa# 
model settings
model = dict(
    type='EncoderDecoder',
    pretrained=checkpoint, # 加载的预训练模型,这里为 Google Research提供的由 JAX 训练框架得到的 Vision Transformer    
    backbone=dict(        
        type='VisionTransformer', # 骨架类名,后面的参数都是该类的初始化参数
        img_size=(512, 512),        
        patch_size=16,        
        in_channels=3,        
        embed_dims=768,        
        num_layers=12,        
        num_heads=12,        
        drop_path_rate=0.1,        
        attn_drop_rate=0.0,        
        drop_rate=0.0,        
        final_norm=True,        
        norm_cfg= dict(type='LN', eps=1e-6, requires_grad=True),   
        with_cls_token=True,        
        interpolate_mode='bicubic',    
    ),

同 OpenMMLab 其他算法库一样,我们使用了 MMCV 中的模块注册机制,通过修改配置文件的 type ,可以使用在 MMSegmentation 已经实现的 backbone 模型。此外,还可以使用 MMClassification 里面的更多主干网络,如 ShuffleNet、EfficientNet 等,可根据 ./configs/convnext 里面 ConvNeXt 的实现方式,详细的方式可以参考: OpenMMLab 模型大联动,MMDet 也能用 MMCls 的网络MMCV 模块注册机制网址

Neck

Neck 可以认为是 backbone 和 head 的连接层,主要负责对 backbone 的特征进行高效融合和增强,能够对输入的单尺度或者多尺度特征进行融合、增强输出等。具体见文件 ./mmseg/models/necks/,v0.24.1 已经实现的 neck 如下:

bdef66cf983fcc313be77d2b605d565b.png 最常用的应该是 FPN,一个典型用法是 ./configs/_base_/models/pointrend_r50.py 里面:

neck=dict(    
    type='FPN',    
    in_channels=[256, 512, 1024, 2048], # 骨架多尺度特征图输出通道    
    out_channels=256, # 增强后通道输出    
    num_outs=4), # 输出num_outs个多尺度特征图

Head

MMSegmentation 的 head 是用来处理 backbone 或 neck 的特征图,对图像里的每个像素 pixel 做分类然后得到分类的结果。具体见文件 ./mmseg/models/decode_heads/,v0.24.1 已经实现的 head 如下: 15f5f597a58b5ae7ecacb449d6a86ab5.png 虽然它们都是用来解码特征图中的信息。但在使用上,可以将它们分为 decoder head,auxiliary head 以及 cascade decoder head。

  • decoder head 是直接在训练和推理中作为图像预测输出的 head。
  • auxiliary head 是只在训练过程中输出图像预测用来辅助损失函数计算的 head。
  • cascade decoder head 是指级联式的 2 个或多个解码头,前一个解码头的输出作为后一个解码头的输入, OCRNet 和 PointRend 两种算法就使用了 cascade decoder head。

在 MMSegmentation 里每个 head 自己单独计算损失,所以把这个公共的行为抽象成了一个基类: BaseDecodeHead,每个算法的 head 都继承自这个基类,类里面包括了计算 loss 的函数,用于计算 head 输出的 logits 值和 label 的损失。

Loss

MMSegmentation 里的 loss 计算的是每个像素上的 logits 和分割标签之间的差别,使用最多的是 cross entropy loss 和 dice loss,v0.24.1 已经实现的 loss 如下: 2d3c0eee5b16bde263943d7cd5ae6a8a.png 除了 ./mmseg/models/losses/ 里的这些 loss 外,计算 loss 时还可以用到一些策略和方法,比如:在线难样本挖掘策略 (OHEM, Online Hard Example Mining) 。

MMSegmentation 跑语义分割数据集

2022.06.07

数据集

无论是MMSegmentation已经支持的数据集,还是自己的数据集,都需要在配置文件里配置数据相关的信息,如数据集本地存储路径,数据预处理流程 Pipeline;继承数据集基类 CustomDataset 以方便调用在某个数据集上加载图像和标注,解析加载数据,评估模型表现等各种功能

数据配置文件

MMSegmentation 的数据集配置基文件在 ./configs/base/datasets 里面,每个数据集配置文件主要包括:(1)data 字段,主要包括dataloader 的配置,例如模型训练时每个 GPU 上面的样本数目和进程数;(2)数据集和数据预处理配置,例如数据集路径和数据预处理 Pipeline。

数据配置文件的 data 字段

这是数据配置文件的一个样例:


data = dict(    
    samples_per_gpu=4,    
    workers_per_gpu=4,    
    train=dict(        
        type='ADE20KDataset',        
        data_root='data/ade/ADEChallengeData2016',        
        img_dir='images/training',        
        ann_dir='annotations/training',        
        pipeline=train_pipeline),    
    val=dict(        
        type='ADE20KDataset',        
        data_root='data/ade/ADEChallengeData2016',        
        img_dir='images/validation',        
        ann_dir='annotations/validation',        
        pipeline=test_pipeline),    
    test=dict(        
        type='ADE20KDataset',        
        data_root='data/ade/ADEChallengeData2016',        
        img_dir='images/validation',        
        ann_dir='annotations/validation',        
        pipeline=test_pipeline))

data 中重要的是如下几个字段:

  • trainval and test: 构建数据集实例的配置,可以通过 registry&build 机制来构建,分别用于模型的训练、验证和测试。
  • samples_per_gpu: 在模型训练时每个 GPU 加载的样本数,它乘以模型训练时的 GPU 数目就是模型训练时的 batch_size。例如,当使用 8 块 GPU 做分布式训练并且 samples_per_gpu=4,那么batch_size is 8*4=32。如果想定义不同 batch_size 用于验证和测试,需要在版本 >=0.24.1 的 MMSegmentation 中使用 val_dataloader 和 test_dataloaser
  • workers_per_gpu: 数据加载时每个 GPU 使用的子进程(subprocess)数目。0 则意味着主进程加载数据。

需要说明的是,samples_per_gpu 仅用于模型训练,因为目前 MMSegmentation 并不支持 batch 方式的推理,所以验证和测试时 samples_per_gpu=1,即每张 GPU 的样本数都是 1。

MMSegmentation 在 v0.24.1 之前,除了 train、 valtestsamples_per_gpu 和 workers_per_gpudata 中的其他字段必须是 PyTorch 中 dataloader 的输入参数,并且模型训练、验证和测试的 dataloaders 都有着同样的输入参数。在 v0.24.1 之后,尽管上述的参数定义仍然可用,但是会将优先支持使用 train_dataloaderval_dataloader和 test_dataloaser 去分别指定模型训练、验证和测试时 dataloader 所需要的参数。

以下就是一个train_dataloaderval_dataloader和 test_dataloaser 使用不同参数的样例:

data = dict(
    samples_per_gpu=4,
    workers_per_gpu=4,
    shuffle=True,
    train=dict(type='xxx', ...),
    val=dict(type='xxx', ...),
    test=dict(type='xxx', ...),
    # 在验证和测试时使用不同的 batch size
    val_dataloader=dict(samples_per_gpu=1, workers_per_gpu=4, shuffle=False),    
    test_dataloader=dict(samples_per_gpu=1, workers_per_gpu=4, shuffle=False))

假如只有一张 GPU 用于模型的训练和测试,因为整体 dataloader 参数定义的优先级比较低,所以训练的 batch size 是 4 并且数据集将会被 shuffle,验证和测试的 batch size 是 1 并且数据集不会被 shuffle。

在 MMSegmentation v0.24.1 之后,我们更推荐使用专门的 dataloader 设置去替代整体 dataloader 的定义,这样可以让数据配置更加清晰易懂。可以修改为:

data = dict(
    train=dict(type='xxx', ...),
    val=dict(type='xxx', ...),
    test=dict(type='xxx', ...),
    # 使用特定的 dataloader 设置
    train_dataloader=dict(samples_per_gpu=4, workers_per_gpu=4, shuffle=True),    
    val_dataloader=dict(samples_per_gpu=1, workers_per_gpu=4, shuffle=False),    
    test_dataloader=dict(samples_per_gpu=1, workers_per_gpu=4, shuffle=False))

数据集预处理 Pipeline

Pipeline 由一系列数据预处理模块组成,得益于 MM 系列模块化的特性,每个模块也都可以单独配置并通过 registry&build 机制来构建。

2eafb975ebfc862fbb683464db5ccb81.png 上图是 MMSegmentation 典型的训练流程 Pipeline,每个模块都接收字典输入,输出也是字典。按照从左到右的顺序执行,绿色表示该模块运行后的新增字段,橙色表示该模块运行后被修改的字段。Pipeline 中的数据变换可以被划分如下:

  • 图片和标签加载,例如 LoadImageFromFile 和 LoadAnnotations
  • 数据处理,例如RandomFlipPhotoMetricDistortion 和 Resize 等,这部分是通常在训练流程中使用。
  • 数据收集,例如 Collect,它会重新收集数据处理的字典,用来准备输入到模型里的数据

在 MMSegmentation 框架中,图片和标签加载和数据处理流程一般是固定的,用户在定制自己数据集的时候,也需要结合具体情况构建合适的 Pipeline。以 ADE20K 数据集为例,在配置文件中,训练时的 Pipeline 如下所示:

# dataset settings
dataset_type = 'ADE20KDataset' # 数据集类型,这将被用来定义数据集。
data_root = 'data/ade/ADEChallengeData2016' # 数据的根路径。
img_norm_cfg = dict( # 图像归一化配置,用来归一化输入的图像。    
    mean=[123.675, 116.28, 103.53], # 预训练里用于预训练主干网络模型的平均值。    
    std=[58.395, 57.12, 57.375], # 预训练里用于预训练主干网络模型的标准差。    
    to_rgb=True) # 预训练里用于预训练主干网络的图像的通道顺序。
crop_size = (512, 512) # 训练时的裁剪大小
train_pipeline = [    
    dict(type='LoadImageFromFile'),    
    dict(type='LoadAnnotations', reduce_zero_label=True),    
    dict(type='Resize', # 变化图像和其注释大小的数据增广。        
        img_scale=(2048, 512), # 图像和标注的 resize 尺度        
        ratio_range=(0.5, 2.0)), # 随机 resize 的比例范围。    
    dict(type='RandomCrop', # 随机裁剪当前图像和其注释。        
        crop_size=crop_size, # 随机裁剪图像生成 patch 的大小。        
        cat_max_ratio=0.75), # 单个类别可以填充的最大区域的比例。    
    dict(type='RandomFlip', # 翻转图像和其注释。        
        prob=0.5),  # 翻转图像的概率    
    dict(type='PhotoMetricDistortion'), # 光学上使用一些方法扭曲当前图像。    
    dict(type='Normalize', # 归一化当前图像数据。         
        **img_norm_cfg),     
    dict(type='Pad', # 填充当前图像到指定大小。        
        size=crop_size, # 填充的图像大小。        
        pad_val=0, # 图像的填充值。        
        seg_pad_val=255), # 'gt_semantic_seg'的填充值。    
    dict(type='DefaultFormatBundle'), # 默认格式转换的组合操作。    
    dict(type='Collect', keys=['img', 'gt_semantic_seg']), # 决定数据里哪些键被传递到分割器里的流程。
]

在构建自己的 Pipeline 时,一定要仔细检查是否真正调用了修改的配置文件,因为新增和修改的字典一旦被错误地覆盖或者遗漏,在运行程序时也可能不会报错,使得排查错误变得困难。

CustomDataset介绍

在训练或验证时经常需要获取该数据集的相关信息,比如获取数据集注释的相关信息,评估数据集预测结果的某些评价指标等等。因此,在 ./mmseg/datasets/custom.py 里面把数据集抽象成一个基类 CustomDataset,在基类里面定义了这些基本的函数,以方便被调用。

CustomDataset 里面主要有以下几个函数:

  • load_annotations():加载全部标注文件,返回一个 List, 其中每个元素是一个字典,键分别是filename 和 ann,值里的信息是图片和对应的标注的文件名。

  • get_ann_info(idx):根据输入inx获取对应语义分割标注的文件名。

  • prepare_train_img(idx):获得经过训练数据处理 Pipeline 之后的训练集的图像数组 img ,和对应的元信息img_metasimg_metas 里的内容可见上一章介绍 Pipeline 的示意图。

  • prepare_test_img(idx):获得经过测试数据流 Pipeline 之后的验证集或测试集的图像数组,和其对应的元信息。字典里内容和prepare_train_img()的一样。

  • __getitem__(idx):通过判断当前是否为训练模式来调用prepare_train_img(idx)prepare_test_img(idx)

  • evaluate(): 评估数据集,输入预测的结果,返回这个数据集所需要的一些评价指标。

因为MMSegmentation 数据集都继承自 CustomDataset,所以熟悉它便熟悉了MMSegmentation 其他数据集的加载、解析和评估的流程。

在介绍完数据集配置文件中需要加入的预处理 Pipeline 和数据集需要继承的 CustomDataset 类之后,下面介绍如何处理自己的数据集,以便训练或验证。

数据集准备和定制化

数据集的准备和定制化,具体可以分成以下几步:

  • 数据集准备,推荐在 mmsegmetaion 目录新建路径 data,然后将数据集转换成 MMSegmentation 可用的格式:分别定义好数据集图像和标注的正确文件路径,其中的标注格式为仅包含每个像素对应标签 id 的单通道标注文件,而不是三通道的 RGB 格式。

  • 在 ./mmseg/datasets 里定义该数据集以注册到 DATASETS 里。

  • 在 ./configs/_base_/datasets 里面设置训练与验证时数据集配置的参数,如数据集路径,数据增强策略等。

上述步骤主要改动的文件位置为:

mmsegmentation   
    |
    |- data   
    |     |- my_dataset                 # 转换后的自己的数据集文件   
    |- mmseg   
    |     |- datasets   
    |     |     |- __init__.py          # 在这里加入自己的数据集的类   
    |     |     |- my_dataset.py               ## 定义自己的数据集的类   
    |     |     |- ...   
    |- configs   
    |     |- _base_   
    |     |     |- datasets   
    |     |     |     |- my_dataset_config.py      # 自己的数据集的配置文件   
    |     |     |- ...   
    |     |- ...

接下来我们详细介绍这三步。

数据集准备

在使用模型做训练、验证和推理前,需要将数据集处理成 MMSegmentation 定制化的格式。对于 MMSegmentation 已经支持的数据集,我们在 ./tools/convert_datasets 中提供了数据集的转换脚本,它们会转换这些原始数据集的大小(例如将较大的遥感数据集裁剪成较小的)和内容(例如将 RBG 格式的标注转换成仅包含每个像素对应标签 id 的单通道标注),同时改变图像和标注的格式与文件夹结构。相关命令可参考数据集准备文档,转换后的数据集可以被 MMSegmentation 一键运行。

数据集最终的目录组织如下,需要将图片放到 img_dir 下,对应的分割标注放到 ann_dir 下:

|- data
│   |- my_dataset
│   │   |- img_dir
│   │   │   |- train
│   │   │   │   |- xxx{img_suffix}
│   │   │   │   |- yyy{img_suffix}
│   │   │   │   |- zzz{img_suffix}
│   │   │   |- val
│   │   |- ann_dir
│   │   │   |- train
│   │   │   │   |- xxx{seg_map_suffix}
│   │   │   │   |- yyy{seg_map_suffix}
│   │   │   │   |- zzz{seg_map_suffix}
│   │   │   |- val

其中 {img_suffix} 和 {seg_map_suffix} 是图像和标注的后缀,常用的是 .png 和 .jpg

实现自己的数据集

生成好上述数据格式后,在 ./mmseg/dataset 里实现数据集,使它可以被注册到 MMCV 的  DATASETS 里面然后被模型调用。实现自己的数据集,只需要继承 CustomDataset 这个类,再定义数据集标注的名称、可视化调色盘以及文件夹后缀格式,如下所示:

from .builder import DATASETS
from .custom import CustomDataset

#将 MyDataset 类注册到 DATASETS 里
@DATASETS.register_module()
class MyDataset(CustomDataset):    
    # 数据集标注的各类名称,即 0, 1, 2, 3... 各个类别的对应名称    
    CLASSES = ('label_a', 'label_b', 'label_c', 'label_d',               
               'label_e', ...)    
    # 各类类别的 BGR 三通道值,用于可视化预测结果    
    PALETTE = [[255, 255, 255], [0, 0, 255], [0, 255, 255], [0, 255, 0], 
               [255, 255, 0], ...]
    
    # 图片和对应的标注,这里对应的文件夹下均为 .png 后缀    
    def __init__(self, **kwargs):        
        super(MyDataset, self).__init__(            
            img_suffix='.png',            
            seg_map_suffix='.png',            
            reduce_zero_label=False, # 此时 label 里的 0(上面 CLASSES 里第一个 “label_a”)在计算损失函数和指标时不会被忽略。            
            **kwargs)

在 ./mmseg/dataset/my_dataset.py 里面定义了数据集的分割类别 CLASSES 和对应的 BGR 通道的调色板 PALETTEPALETTE 只在预测结果可视化的时候会用到,并不会影响训练和验证。需要强调的是,如果 label 中的 0 是背景并且想在计算评价指标的时候忽略掉它,需要设置 reduce_zero_label=True
其在 ./mmseg/core/evaluation/metrics.py 中的原理是:当设置 reduce_zero_label=True 时,会修改分割的标签类别,将 index 为 0 的类别安排到 255,所以在训练和和测试加载分割标注时,都会做如下操作:

if reduce_zero_label:    
    label[label == 0] = 255    
    label = label - 1    
    label[label == 254] = 255

255 是标签里被忽略的 index。创建好 ./mmseg/dataset/my_dataset.py 后,需要在 ./mmseg/dataset/__init__.py 里也加入它:

# Copyright (c) OpenMMLab. All rights reserved.
from .my_dataset import MyDataset
__all__ = [
    ...,    
    'MyDataset'
]

设置数据集配置文件

数据集定义好后,还需要在 ./configs/_base_/datasets 里面定义该数据集有关的配置项 my_dataset_config.py,使之与其他的配置参数一起在训练和测试时调用。 首先简单介绍下为何多了一个 _base_ 基配置文件夹:自从 2020 年 6 月 MMDetection 发布 V2.0 版本以来,OpenMMLab 代码库设计了新的 config 系统,支持了多重继承机制。将常用的数据集配置、基础模型以及训练策略放到了 ./configs/_base_/ 文件夹中。每个新的 config 只需要继承一个或者多个已有的 config,然后对其中需要修改的字段进行重载。通过将 config 继承的层级控制在可接受范围内,提升了配置文件的可维护性。

以下为数据集配置文件的一个示例:

# 在./mmseg/datasets/__init__.py 中定义的数据集类型
dataset_type = 'MyDataset'
# 数据集准备生成的文件夹路径
data_root = 'data/my_dataset'

img_norm_cfg = dict( # 常用这组参数归一化是因为它是 ImageNet 1K 预训练使用的图像均值与方差  
    mean=[123.675, 116.28, 103.53], 
    std=[58.395, 57.12, 57.375], 
    to_rgb=True
)
crop_size = (512, 512) # 训练时图像裁剪的大小
train_pipeline = [    
    dict(type='LoadImageFromFile'),    
    dict(type='LoadAnnotations', reduce_zero_label=True),    
    dict(type='Resize', img_scale=(512, 512), ratio_range=(0.5, 2.0)),    
    dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),    
    dict(type='RandomFlip', prob=0.5),    
    dict(type='PhotoMetricDistortion'),    
    dict(type='Normalize', **img_norm_cfg),    
    dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),    
    dict(type='DefaultFormatBundle'),    
    dict(type='Collect', keys=['img', 'gt_semantic_seg']),
]
test_pipeline = [    
    dict(type='LoadImageFromFile'),    
    dict(       
        type='MultiScaleFlipAug',        
        img_scale=(512, 512),        
        # img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],        
        flip=False,        
        transforms=[            
            dict(type='Resize', keep_ratio=True),            
            dict(type='RandomFlip'),            
            dict(type='Normalize', **img_norm_cfg),            
            dict(type='ImageToTensor', keys=['img']),            
            dict(type='Collect', keys=['img']),        
            ])
]
data = dict(    
    samples_per_gpu=4, # 单个 GPU 的 Batch size    
    workers_per_gpu=4, # 单个 GPU 分配的数据加载线程数    
    train=dict( # 训练数据集配置        
        type=dataset_type, # 数据集的类别, 细节参考自 mmseg/datasets/        
        data_root=data_root, # 数据集的根目录。        
        img_dir='img_dir/train', # 数据集图像的文件夹        
        ann_dir='ann_dir/train', # 数据集注释的文件夹        
        pipeline=train_pipeline), # 流程, 由之前创建的 train_pipeline 传递进来    
        val=dict( # 验证数据集的配置        
        type=dataset_type,        
        data_root=data_root,        
        img_dir='img_dir/val',        
        ann_dir='ann_dir/val',        
        pipeline=test_pipeline), # 由之前创建的 test_pipeline 传递的流程    
    test=dict(        
        type=dataset_type,        
        data_root=data_root,        
        img_dir='img_dir/val',        
        ann_dir='ann_dir/val',        
        pipeline=test_pipeline)
)

各个配置项的具体作用可以参考配置文件教程。至此,定义的数据集就完成了数据集的准备和定制化,只需要在 ./configs/ 里创建的配置文件里调用该数据集即可。例如:

_base_ = [    
    '../_base_/models/pspnet_r50-d8.py', '../_base_/datasets/my_dataset_config.py',    
    '../_base_/default_runtime.py', '../_base_/schedules/schedule_80k.py'
]
model = dict(
    decode_head=dict(num_classes=YOUR_DATASET_CLASSES), auxiliary_head=dict(num_classes=YOUR_DATASET_CLASSES))

本文主要参考公众号OpenMMLab的文章《带你轻松掌握 MMSegmentation 整体构建流程》、《# 带你轻松用 MMSegmentation 跑语义分割数据集》