mmsegmentation自定义数据集的准备，配置文件编写以及训练，测试

一、前提

确认已经安装好了mmsegmentation的环境。具体安装方法，请看官方的get_started.md。

二、数据准备

首先是结合官方的customize_datasets.md，如图：

其中img_dir中存放的就是原图，ann_dir中就是放的对应图像的"mask"。

我的数据是用labelme标注的，出来就是一张图对应一个json文件，它是只标注了每个目标物的外轮廓，大概这个样子：

然后接下来就是按照官方的说法，制作每张图对应的ann:

import os
import json
import glob
import shutil
import tqdm
import cv2
import numpy as np
from PIL import Image
from sklearn.model_selection import train_test_split

np.random.seed(0)
w, h = (1280, 1024)   # 这是我的图的大小，你看你来改


def labelme2seg(json_files: list, output_path: str):
    for json_file in tqdm.tqdm(json_files, desc="transforming："):
        img_data = np.ones((h, w), dtype=np.uint8) * 27   # 我一共27类，你目标物多少类这里就写多少类，
        with open(json_file, encoding="utf-8") as f:
            json_data = json.load(f)
        labels_data = json_data["shapes"]
        # 将目标物区域像素填充为对应ID号
        for label_data in labels_data:
            # 下面这行，你的label不是数字的话，是汉字或者其它，自己记得稍微改一下，映射成数字，从0开始
            goods_id = int(label_data["label"])
            location = np.asarray(label_data["points"], dtype=np.int32)
            cv2.fillPoly(img_data, [location], color=(goods_id, goods_id, goods_id))

        res_img = Image.fromarray(img_data, mode="P")
        res_img_name = os.path.basename(json_file).replace(".json", ".png")
        res_img.save(os.path.join(output_path, res_img_name))
    return


if __name__ == '__main__':
    labelme_path = r"上面方.jpg和json文件的路径"
    save_path = r"最终保存的路径/my_dataset"

    img_dir_train = os.path.join(save_path, "img_dir", "train")
    img_dir_val = os.path.join(save_path, "img_dir", "val")
    img_dir_test = os.path.join(save_path, "img_dir", "test")

    ann_dit_train = os.path.join(save_path, "ann_dir", "train")
    ann_dir_val = os.path.join(save_path, "ann_dir", "val")
    ann_dir_test = os.path.join(save_path, "ann_dir", "test")

    if not os.path.exists(img_dir_train):
        os.makedirs(img_dir_train)
    if not os.path.exists(img_dir_val):
        os.makedirs(img_dir_val)
    if not os.path.exists(img_dir_test):
        os.makedirs(img_dir_test)

    if not os.path.exists(ann_dit_train):
        os.makedirs(ann_dit_train)
    if not os.path.exists(ann_dir_val):
        os.makedirs(ann_dir_val)
    if not os.path.exists(ann_dir_test):
        os.makedirs(ann_dir_test)

    json_list_path = glob.glob(labelme_path + "/*.json")
    train_path, test_val_path = train_test_split(json_list_path, test_size=0.15)
    test_path, val_path = train_test_split(test_val_path, test_size=0.15)

    # 制作mask：
    labelme2seg(train_path, ann_dit_train)
    labelme2seg(val_path, ann_dir_val)
    labelme2seg(test_path, ann_dir_test)

    # 图复制进对应位置
    for file in tqdm.tqdm(train_path, desc="copy train_img"):
        shutil.copy(file.replace(".json", ".jpg"), img_dir_train)
    for file in tqdm.tqdm(val_path, desc="copy val_img"):
        shutil.copy(file.replace(".json", ".jpg"), img_dir_val)
    for file in tqdm.tqdm(test_path, desc="copy test_img"):
        shutil.copy(file.replace(".json", ".jpg"), img_dir_test)

然后等你制作好了数据集，就把这个my_dataset文件夹放在mmsegmentation的data目录下。
特别注意：一定要看以上代码，我写的3处注释，根据你的实际情况来改，特别是类别名映射成[0, class_num-1]。除非极度巧合，不然代码肯定的是会报错的。

三、配置文件

自定义数据类

在mmseg/datasets新建一个文件my_dataset.py仿造其它的数据集写下如下内容：

from .builder import DATASETS
from .custom import CustomDataset


@DATASETS.register_module()
class MyDataset(CustomDataset):
    # 写你实际的类别名就好了，最后再加上一个background
    CLASSES = (
        'T01', 'T02', 'T03', 'T04', 'T05', 'T06', 'T07', 'T08', 'T09', 'T10', 'T11', 'T12', 'T13',
        'T14', 'T15', 'T16', 'T17', 'T18', 'T19', 'T20', 'T21', 'T22', 'T23', 'T24', 'T25','T26', 'background'
    )
    # 这个数量与上面个数对应就好了
    PALETTE = [[120, 120, 120], [180, 120, 120], [6, 230, 230], [80, 50, 50],
               [4, 200, 3], [120, 120, 80], [140, 140, 140], [204, 5, 255],
               [230, 230, 230], [4, 250, 7], [224, 5, 255], [235, 255, 7],
               [150, 5, 61], [120, 120, 70], [8, 255, 51], [255, 6, 82],
               [143, 255, 140], [204, 255, 4], [255, 51, 7], [204, 70, 3],
               [0, 102, 200], [61, 230, 250], [255, 6, 51], [11, 102, 255],
               [255, 7, 71], [255, 9, 224], [9, 7, 230], [220, 220, 220]]

    def __init__(self, **kwargs):
        super(MyDataset, self).__init__(
            **kwargs
        )

然后在mmseg/datasets/__init__.py中把自己的数据集添加进去(主要是添加以下两行)：

from .my_dataset import MyDataset
# 在 __all__中添加自己的类名
__all__ = [
    '......', 'LoveDADataset', 'MyDataset'       # 最后添加这个
]

自定义配置文件

首先说明，这是仿照configs/pspnet/pspnet_r50-d8_512x512_160k_ade20k.py这一配置文件来写的，其它的大同小异，把pspnet_r50-d8_512x512_160k_ade20k.py配置文件中的是用到的几个配置文件都写进了一个文件里my_psp.py，放在configs/pspnet下，内容如下：

# 参考 pspnet_r50-d8_512x512_160k_ade20k 来做的

# 1.model：根据'../_base_/models/pspnet_r50-d8.py'
# model settings
norm_cfg = dict(type='SyncBN', requires_grad=True)
model = dict(
    type='EncoderDecoder',
    pretrained='open-mmlab://resnet50_v1c',
    backbone=dict(
        type='ResNetV1c',
        depth=50,
        num_stages=4,
        out_indices=(0, 1, 2, 3),
        dilations=(1, 1, 2, 4),
        strides=(1, 2, 1, 1),
        norm_cfg=norm_cfg,
        norm_eval=False,
        style='pytorch',
        contract_dilation=True),
    decode_head=dict(
        type='PSPHead',
        in_channels=2048,
        in_index=3,
        channels=512,
        pool_scales=(1, 2, 3, 6),
        dropout_ratio=0.1,
        num_classes=28,   # 注意改这个类别数,我的是27类+最后一个background
        norm_cfg=norm_cfg,
        align_corners=False,
        loss_decode=dict(
            type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)),
    auxiliary_head=dict(
        type='FCNHead',
        in_channels=1024,
        in_index=2,
        channels=256,
        num_convs=1,
        concat_input=False,
        dropout_ratio=0.1,
        num_classes=28,       # 注意改这个值，
        norm_cfg=norm_cfg,
        align_corners=False,
        loss_decode=dict(
            type='CrossEntropyLoss', use_sigmoid=False, loss_weight=0.4)),
    # model training and testing settings
    train_cfg=dict(),
    test_cfg=dict(mode='whole'))


# 2. data：根据'../_base_/datasets/ade20k.py'
# dataset settings
dataset_type = 'MyDataset'   # 自己mmseg/datasets/my_dataset.py中的类名
data_root = 'data/my_dataset'
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=(1280, 1024), ratio_range=(0.5, 2.0)),   # img_scale我还是就用的我原图的大小，你可以改成你的大小，其他地方同步改
    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=(1280, 1024),
        # 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,   # 显存不够，来把这改小
    workers_per_gpu=4,    
    train=dict(
        type=dataset_type,
        data_root=data_root,
        img_dir='img_dir/train',
        ann_dir='ann_dir/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=dict(
        type=dataset_type,
        data_root=data_root,
        img_dir='img_dir/test',
        ann_dir='ann_dir/test',
        pipeline=test_pipeline))


# 3. 根据：'../_base_/default_runtime.py'
# yapf:disable
log_config = dict(
    interval=50,
    hooks=[
        dict(type='TextLoggerHook', by_epoch=False),
        # dict(type='TensorboardLoggerHook')
    ])
# yapf:enable
dist_params = dict(backend='nccl')
log_level = 'INFO'
load_from = None
resume_from = None
workflow = [('train', 1)]
cudnn_benchmark = True


# 4. schedules：根据'../_base_/schedules/schedule_160k.py'
# optimizer
optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005)   # 根据自己情况改学习率吧
optimizer_config = dict()
# learning policy
lr_config = dict(policy='poly', power=0.9, min_lr=1e-4, by_epoch=False)
# runtime settings
runner = dict(type='IterBasedRunner', max_iters=160000)
checkpoint_config = dict(by_epoch=False, interval=16000)
evaluation = dict(interval=16000, metric='mIoU', pre_eval=True)

训练

在mmsegmentation目录下，激活你的虚拟环境，然后执行：

单张GPU：python tools/train.py configs/pspnet/my_psp.py

两张GPU：./tools/dist_train.sh configs/pspnet/my_psp.py 2

然后你应该就能看到这个效果：

然后就大功告成。

四、结果测试

把测试集的结果保存起来：

python tools/test.py configs/pspnet/my_psp.py work_dirs/my_psp/latest.pth --show-dir ./mypsp_output

如果想像官方demo展示一样的话，就还需要修改文件mmseg/core/evaluation/class_names.py:

仿照里面已有的数据写，添加（与上面保持一致）

# for myself dataset
def mydata_classes():    
    return [
        'T01', 'T02_1', 'T02_2', 'T02_3', 'T03_1', 'T03_2', 'T03_3', 'T04', 'T05', 'T06', 'T07', 'T08', 'T09', 'T10',
        'T11', 'T12', 'T13', 'T14', 'T15', 'T16', 'T17', 'T18', 'T19', 'T20', 'T21', 'T22', 'T26', 'background'
    ]

def mydata_palette():
    return [[120, 120, 120], [180, 120, 120], [6, 230, 230], [80, 50, 50],
            [4, 200, 3], [120, 120, 80], [140, 140, 140], [204, 5, 255],
            [230, 230, 230], [4, 250, 7], [224, 5, 255], [235, 255, 7],
            [150, 5, 61], [120, 120, 70], [8, 255, 51], [255, 6, 82],
            [143, 255, 140], [204, 255, 4], [255, 51, 7], [204, 70, 3],
            [0, 102, 200], [61, 230, 250], [255, 6, 51], [11, 102, 255],
            [255, 7, 71], [255, 9, 224], [9, 7, 230], [220, 220, 220]]

这里面还有一处需要修改：

dataset_aliases = {
    'cityscapes': ['cityscapes'],
    'ade': ['ade', 'ade20k'],
    'voc': ['voc', 'pascal_voc', 'voc12', 'voc12aug'],
    'mydata': ['mydata']     # 这处是自己添加的
}

Tips：函数名这些别写错了，想要改成别的也一定要全部一起改。

然后可以用官方展示demo的方式来展示你的结果了：

python demo/image_demo.py data/my_dataset/img_dir/test/10740.jpg work_dirs/my_psp/my_psp.py work_dirs/my_psp/latest.pth --device cuda:0 --palette mydata

再次说明，一定要看我代码里写的注释说明，根据你的实际情况来改，除非极度巧合，不然直接拿来就跑，肯定是会出错的。

我这里都是自己成功跑通了的，也是学习摸索中，如果哪里有什么问题，还请交流、雅正。

希望能帮到你~