MMSegmentation笔记

如何训练自制数据集？

首先需要在 mmsegmentation/mmseg/datasets 目录下创建一个自制数据集的配置文件，以我的苹果叶片病害分割数据集为例，创建了mmsegmentation/mmseg/datasets/appleleafseg.py

可以看到，这个配置文件主要定义了自制数据集中的 METAINFO , 包括标签的类别，以及对应的 palette 调色板色彩数值，还定义了原始图像和标签图像的文件后缀，分别是 jpg 和 png，以及设置 reduce_zero_label 属性 (是否忽略背景)

from mmseg.registry import DATASETS
from .basesegdataset import BaseSegDataset


@DATASETS.register_module()
class AppleLeafSegDataset(BaseSegDataset):
    METAINFO = dict(
        classes=('background', 'Alternaria_Boltch', 'Brown_spot', 'Frogeye_leaf_spot', 'Grey_spot', 'Mosaic', 'Powdery_mildew', 'Rust', 'Scab', 'Health'),
        palette=[[0, 0, 0], [170, 0, 0], [99, 102, 129], [249, 193, 0], [160, 180, 0],
                 [115, 82, 59], [217, 213, 180], [51, 142, 137], [218, 147, 70], [234, 132, 163]])
    def __init__(self,
                 img_suffix='.jpg',
                 seg_map_suffix='.png',
                 reduce_zero_label=False,# 因为上面METAINFO已经将背景0作为一种类别并且设置掩码色彩为0,0,0所以这里的reduce_zero_label需要设置为false
                 **kwargs) -> None:
        super().__init__(
            img_suffix=img_suffix,
            seg_map_suffix=seg_map_suffix,
            reduce_zero_label=reduce_zero_label,
            **kwargs)

然后将 AppleLeafSegDataset 添加到 mmseg/datasets/__init__.py 中的__all__里

__all__ = [
    'BaseSegDataset', 'BioMedical3DRandomCrop', 'BioMedical3DRandomFlip',
    'CityscapesDataset', 'PascalVOCDataset', 'ADE20KDataset',
    'PascalContextDataset', 'PascalContextDataset59', 'ChaseDB1Dataset',
    'DRIVEDataset', 'HRFDataset', 'STAREDataset', 'DarkZurichDataset',
    'NightDrivingDataset', 'COCOStuffDataset', 'LoveDADataset',
    'MultiImageMixDataset', 'iSAIDDataset', 'ISPRSDataset', 'PotsdamDataset',
    'LoadAnnotations', 'RandomCrop', 'SegRescale', 'PhotoMetricDistortion',
    'RandomRotate', 'AdjustGamma', 'CLAHE', 'Rerange', 'RGB2Gray',
    'RandomCutOut', 'RandomMosaic', 'PackSegInputs', 'ResizeToMultiple',
    'LoadImageFromNDArray', 'LoadBiomedicalImageFromFile',
    'LoadBiomedicalAnnotation', 'LoadBiomedicalData', 'GenerateEdge',
    'DecathlonDataset', 'LIPDataset', 'ResizeShortestEdge',
    'BioMedicalGaussianNoise', 'BioMedicalGaussianBlur',
    'BioMedicalRandomGamma', 'BioMedical3DPad', 'RandomRotFlip',
    'SynapseDataset', 'REFUGEDataset', 'MapillaryDataset_v1',
    'MapillaryDataset_v2', 'Albu', 'LEVIRCDDataset',
    'LoadMultipleRSImageFromFile', 'LoadSingleRSImageFromFile',
    'ConcatCDInput', 'BaseCDDataset', 'DSDLSegDataset', 'BDD100KDataset',
    'NYUDataset', 'HSIDrive20Dataset', 'AppleLeafSegDataset'
]

接下来，需要在 mmsegmentation/mmseg/utils/class_names.py 中补充数据集元信息

我的苹果树叶病害数据集相关片段如下：

def appleleafdiseases_classes():
    """BDD100K class names for external use(the class name is compatible with
    Cityscapes )."""
    return [
        'background', 'Alternaria_Boltch', 'Brown_spot', 
        'Frogeye_leaf_spot', 'Grey_spot', 'Mosaic',
        'Powdery_mildew', 'Rust', 'Scab', 'Health'
    ]


def appleleafdiseases_palette():
    """bdd100k palette for external use(same with cityscapes)"""
    return [[0, 0, 0], [170, 0, 0], [99, 102, 129], [249, 193, 0], [160, 180, 0],
            [115, 82, 59], [217, 213, 180], [51, 142, 137], [218, 147, 70], [234, 132, 163]]
            
dataset_aliases = {
    'cityscapes': ['cityscapes'],
    'ade': ['ade', 'ade20k'],
    'voc': ['voc', 'pascal_voc', 'voc12', 'voc12aug'],
    'pcontext': ['pcontext', 'pascal_context', 'voc2010'],
    'loveda': ['loveda'],
    'potsdam': ['potsdam'],
    'vaihingen': ['vaihingen'],
    'cocostuff': [
        'cocostuff', 'cocostuff10k', 'cocostuff164k', 'coco-stuff',
        'coco-stuff10k', 'coco-stuff164k', 'coco_stuff', 'coco_stuff10k',
        'coco_stuff164k'
    ],
    'isaid': ['isaid', 'iSAID'],
    'stare': ['stare', 'STARE'],
    'lip': ['LIP', 'lip'],
    'mapillary_v1': ['mapillary_v1'],
    'mapillary_v2': ['mapillary_v2'],
    'bdd100k': ['bdd100k'],
    'hsidrive': [
        'hsidrive', 'HSIDrive', 'HSI-Drive', 'hsidrive20', 'HSIDrive20',
        'HSI-Drive20'
    ],
    'appleleafdiseases': ['appleleafdiseases']
}

然后，需要在mmsegmentation/configs/_base_/datasets/目录下创建一个新的数据集配置文件 mmsegmentation/configs/_base_/datasets/apple.py

这个数据集配置文件代码如下，可以看到，主要是告诉模型训练和测试的一些配置信息，包括数据集类和数据集路径，训练，测试的pipiline数据增强，不同的dataloader(训练集，验证集，测试集)，验证集测试集的评价指标计算。

# dataset settings
dataset_type = 'AppleLeafSegDataset'
data_root = 'AppleLeafSegDataset/' # 自己数据集所在位置
img_scale = (320, 640) # img_scale是指图像在处理管道中将被调整到的尺寸
crop_size = (160, 320)
train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations', reduce_zero_label=False), # 不忽略背景
    dict(
        type='RandomResize',
        scale=img_scale,
        ratio_range=(0.5, 2.0),
        keep_ratio=True),
    dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
    dict(type='RandomFlip', prob=0.5),
    dict(type='PhotoMetricDistortion'),
    dict(type='PackSegInputs')
]
test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='Resize', scale=img_scale, keep_ratio=True),
    # add loading annotation after ``Resize`` because ground truth
    # does not need to do resize data transform
    dict(type='LoadAnnotations', reduce_zero_label=False),
    dict(type='PackSegInputs')
]
img_ratios = [0.5, 0.75, 1.0, 1.25, 1.5, 1.75]
# 测试时增强 (TTA) 是一种在测试阶段使用的数据增强策略。它对同一张图片应用不同的增强，例如翻转和缩放，用于模型推理，然后将每个增强后的图像的预测结果合并，以获得更准确的预测结果。
tta_pipeline = [
    dict(type='LoadImageFromFile', backend_args=None),
    dict(
        type='TestTimeAug',
        transforms=[
            [
                dict(type='Resize', scale_factor=r, keep_ratio=True)
                for r in img_ratios
            ],
            [
                dict(type='RandomFlip', prob=0., direction='horizontal'),
                dict(type='RandomFlip', prob=1., direction='horizontal')
            ], [dict(type='LoadAnnotations')], [dict(type='PackSegInputs')]
        ])
]
train_dataloader = dict(
    batch_size=4,
    num_workers=4,
    persistent_workers=True,
    sampler=dict(type='InfiniteSampler', shuffle=True),
    dataset=dict(
        type=dataset_type,
        data_root=data_root,
        data_prefix=dict(
            img_path='images/training', seg_map_path='annotations/training'),
        pipeline=train_pipeline))
val_dataloader = dict(
    batch_size=1,
    num_workers=4,
    persistent_workers=True,
    sampler=dict(type='DefaultSampler', shuffle=False),
    dataset=dict(
        type=dataset_type,
        data_root=data_root,
        data_prefix=dict(
            img_path='images/validation',
            seg_map_path='annotations/validation'),
        pipeline=test_pipeline))
test_dataloader = val_dataloader

val_evaluator = dict(type='IoUMetric', iou_metrics=['mIoU'])
test_evaluator = val_evaluator

最后，我们需要创建一个总的配置文件，mmsegmentation/configs/unet/unet_s5-d16_deeplabv3_4xb4-40k_appleleafdiseases-320×640.py

这里可以选择mmsegmentation/configs/目录下的不同模型进行实验，这里以unet为例，我创建的这个文件代码如下：
可以看到，_base_定义了模型配置，数据集配置，调度策略配置，运行时配置。
然后也定义了裁剪大小，数据预处理。

_base_ = [
    '../_base_/models/apple_deeplabv3_unet_s5-d16.py', '../_base_/datasets/apple.py',
    '../_base_/default_runtime.py', '../_base_/schedules/schedule_40k.py'
]
crop_size = (160, 320)
data_preprocessor = dict(size=crop_size)
model = dict(
    data_preprocessor=data_preprocessor,
    test_cfg=dict(crop_size=(160, 320), stride=(85, 85)))

然后，创建一个mmsegmentation/configs/_base_/models/apple_deeplabv3_unet_s5-d16.py

代码如下, 可以看到定义了数据预处理，模型结构，backbone类型，解码器头和辅助解码器头：

# model settings
norm_cfg = dict(type='BN', requires_grad=True)
data_preprocessor = dict(
    type='SegDataPreProcessor',
    mean=[123.675, 116.28, 103.53],
    std=[58.395, 57.12, 57.375],
    bgr_to_rgb=True,
    pad_val=0,
    seg_pad_val=255)
model = dict(
    type='EncoderDecoder',
    data_preprocessor=data_preprocessor,
    pretrained=None,
    backbone=dict(
        type='UNet',
        in_channels=3,
        base_channels=64,
        num_stages=5,
        strides=(1, 1, 1, 1, 1),
        enc_num_convs=(2, 2, 2, 2, 2),
        dec_num_convs=(2, 2, 2, 2),
        downsamples=(True, True, True, True),
        enc_dilations=(1, 1, 1, 1, 1),
        dec_dilations=(1, 1, 1, 1),
        with_cp=False,
        conv_cfg=None,
        norm_cfg=norm_cfg,
        act_cfg=dict(type='ReLU'),
        upsample_cfg=dict(type='InterpConv'),
        norm_eval=False),
    decode_head=dict(
        type='ASPPHead',
        in_channels=64,
        in_index=4,
        channels=16,
        dilations=(1, 12, 24, 36),
        dropout_ratio=0.1,
        num_classes=10,
        norm_cfg=norm_cfg,
        align_corners=False,
        loss_decode=dict(
            type='LovaszLoss', reduction='none', loss_weight=1.0)),
    auxiliary_head=dict(
        type='FCNHead',
        in_channels=128,
        in_index=3,
        channels=64,
        num_convs=1,
        concat_input=False,
        dropout_ratio=0.1,
        num_classes=10,
        norm_cfg=norm_cfg,
        align_corners=False,
        loss_decode=dict(
            type='LovaszLoss', reduction='none', loss_weight=0.4)),
    # model training and testing settings
    train_cfg=dict(),
    test_cfg=dict(mode='slide', crop_size=128, stride=85))

然后，重新启动

python setup.py install
pip install -v -e .

开始训练

python tools/train.py configs/unet/unet_s5-d16_deeplabv3_4xb4-40k_appleleafdiseases-320×640.py --work-dir mmseg_log