Mask R-CNN训练自己的数据集

数据集制作

通常使用labelme来制作实例分割数据集，也有教程和代码来转换成COCO数据集。labelme项目地址为：https://github.com/wkentaro/labelme/tree/main

安装labelme

conda create --name=labelme python=3
conda activate labelme
pip install labelme

# or install standalone executable/app from:
# https://github.com/wkentaro/labelme/releases
1
2
3
4
5
6

标注分割区域

在labelme标注区域时，对于存在遮挡的物体，可以利用labelme标签里的group选项。如下图所示，elephant有两部分区域，group都设置为0.

转换为COCO数据集

在labelme项目下的examples/instance_segmentation文件夹中提供转VOC和COCO两种格式的数据和脚本。本文只对转COCO格式进行描述，文件结构如下所示。



对于自定义数据集，按照以上的结果准备好图像数据和标签数据，即data_annotated文件夹中的内容。运行如下代码，转换为COCO格式的数据集。

python labelme2coco.py data_annotated/ coco --labels labels.txt
1

完成之后，会在输出文件夹下得到如下的内容。

一个小的点，在保存json文件时，可以将代码修改成如下，得到的json文件看起来比较美观，同时支持中文。

with open(out_ann_file, "w") as f:
    json.dump(data, f, indent=2, ensure_ascii=False))
    #ensure_ascii=False可以消除json包含中文的乱码问题
1
2
3

Mask R-CNN训练

本文的环境配置如下：

• pytorch==1.7.0
• torchvision==0.8.0
• mmcv-full==1.2.7
• mmdet==2.8.0

config文件修改

model config

在model的配置部分，唯一需要修改的是num_classes参数，根据数据集修改对应值。

# model settings

num_classes=1

model = dict(
    type='MaskRCNN',
    pretrained='torchvision://resnet50',
    backbone=dict(
        type='ResNet',
        depth=50,
        num_stages=4,
        out_indices=(0, 1, 2, 3),
        frozen_stages=1,
        norm_cfg=dict(type='BN', requires_grad=True),
        norm_eval=True,
        style='pytorch'),
    neck=dict(
        type='FPN',
        in_channels=[256, 512, 1024, 2048],
        out_channels=256,
        num_outs=5),
    rpn_head=dict(
        type='RPNHead',
        in_channels=256,
        feat_channels=256,
        anchor_generator=dict(
            type='AnchorGenerator',
            scales=[8],
            ratios=[0.5, 1.0, 2.0],
            strides=[4, 8, 16, 32, 64]),
        bbox_coder=dict(
            type='DeltaXYWHBBoxCoder',
            target_means=[.0, .0, .0, .0],
            target_stds=[1.0, 1.0, 1.0, 1.0]),
        loss_cls=dict(
            type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0),
        loss_bbox=dict(type='L1Loss', loss_weight=1.0)),
    roi_head=dict(
        type='StandardRoIHead',
        bbox_roi_extractor=dict(
            type='SingleRoIExtractor',
            roi_layer=dict(type='RoIAlign', output_size=7, sampling_ratio=0),
            out_channels=256,
            featmap_strides=[4, 8, 16, 32]),
        bbox_head=dict(
            type='Shared2FCBBoxHead',
            in_channels=256,
            fc_out_channels=1024,
            roi_feat_size=7,
            num_classes=num_classes,
            bbox_coder=dict(
                type='DeltaXYWHBBoxCoder',
                target_means=[0., 0., 0., 0.],
                target_stds=[0.1, 0.1, 0.2, 0.2]),
            reg_class_agnostic=False,
            loss_cls=dict(
                type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0),
            loss_bbox=dict(type='L1Loss', loss_weight=1.0)),
        mask_roi_extractor=dict(
            type='SingleRoIExtractor',
            roi_layer=dict(type='RoIAlign', output_size=14, sampling_ratio=0),
            out_channels=256,
            featmap_strides=[4, 8, 16, 32]),
        mask_head=dict(
            type='FCNMaskHead',
            num_convs=4,
            in_channels=256,
            conv_out_channels=256,
            num_classes=num_classes,
            loss_mask=dict(
                type='CrossEntropyLoss', use_mask=True, loss_weight=1.0))))
# model training and testing settings
train_cfg = dict(
    rpn=dict(
        assigner=dict(
            type='MaxIoUAssigner',
            pos_iou_thr=0.7,
            neg_iou_thr=0.3,
            min_pos_iou=0.3,
            match_low_quality=True,
            ignore_iof_thr=-1),
        sampler=dict(
            type='RandomSampler',
            num=256,
            pos_fraction=0.5,
            neg_pos_ub=-1,
            add_gt_as_proposals=False),
        allowed_border=-1,
        pos_weight=-1,
        debug=False),
    rpn_proposal=dict(
        nms_across_levels=False,
        nms_pre=2000,
        nms_post=1000,
        max_num=1000,
        nms_thr=0.7,
        min_bbox_size=0),
    rcnn=dict(
        assigner=dict(
            type='MaxIoUAssigner',
            pos_iou_thr=0.5,
            neg_iou_thr=0.5,
            min_pos_iou=0.5,
            match_low_quality=True,
            ignore_iof_thr=-1),
        sampler=dict(
            type='RandomSampler',
            num=512,
            pos_fraction=0.25,
            neg_pos_ub=-1,
            add_gt_as_proposals=True),
        mask_size=28,
        pos_weight=-1,
        debug=False))
test_cfg = dict(
    rpn=dict(
        nms_across_levels=False,
        nms_pre=1000,
        nms_post=1000,
        max_num=1000,
        nms_thr=0.7,
        min_bbox_size=0),
    rcnn=dict(
        score_thr=0.05,
        nms=dict(type='nms', iou_threshold=0.5),
        max_per_img=100,
        mask_thr_binary=0.5))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127

data config

在data的配置部分，需要修改data_root，classes参数来指明数据集的路径，以及对应的类别名列表。对于训练集、验证集和测试集的ann_file和img_prefix两个参数需要进行调整。

dataset_type = 'CocoDataset'

img_norm_cfg = dict(
    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)

train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations', with_bbox=True, with_mask=True),
    dict(type='Resize', img_scale=(416, 416), keep_ratio=True),
    dict(type='RandomFlip', flip_ratio=0.5),
    dict(type='Normalize', **img_norm_cfg),
    dict(type='Pad', size_divisor=32),
    dict(type='DefaultFormatBundle'),
    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels', 'gt_masks']),
]
test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(
        type='MultiScaleFlipAug',
        img_scale=(416, 416),
        flip=False,
        transforms=[
            dict(type='Resize', keep_ratio=True),
            dict(type='RandomFlip'),
            dict(type='Normalize', **img_norm_cfg),
            dict(type='Pad', size_divisor=32),
            dict(type='ImageToTensor', keys=['img']),
            dict(type='Collect', keys=['img']),
        ])
]

data_root = 'datasets/xuzhou2_single_jietou/'
classes=["jietou"]
data = dict(
    samples_per_gpu=32,
    workers_per_gpu=1,
    # dataset type
    train=dict(
        type=dataset_type,
        ann_file=data_root + 'annotations/instances_jietou_train20231016.json',
        img_prefix=data_root + 'train/',
        pipeline=train_pipeline,
        classes=classes
        ),
    val=dict(
        type=dataset_type,
        ann_file=data_root + 'annotations/instances_jietou_val20231016.json',
        img_prefix=data_root + 'val/',
        pipeline=test_pipeline,
        classes=classes
        ),
    test=dict(
        type=dataset_type,
        ann_file=data_root + 'annotations/instances_jietou_val20231016.json',
        img_prefix=data_root + 'val/',
        pipeline=test_pipeline,
        classes=classes
        ),
    )
evaluation = dict(
                    interval=10,   
                    metric=['bbox', 'segm']
                    )
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63

优化器和学习率的配置

使用随机梯度下降法来更新参数，修改学习率的优化策略为warmup+余弦衰减策略。

# optimizer
optimizer = dict(type='SGD', lr=0.02, momentum=0.9, weight_decay=0.0001)
optimizer_config = dict(grad_clip=None)

# Learning rate scheduler config used to register LrUpdater hook
lr_config = dict(
    policy='CosineAnnealing', 
    min_lr=0,
    warmup='linear',
    warmup_iters=25,
    warmup_ratio=0.001,
    warmup_by_epoch=True
)
total_epochs = 150
1
2
3
4
5
6
7
8
9
10
11
12
13
14

runtime配置

修改权重保存间隔为5个epoch保存一次。

checkpoint_config = dict(interval=5)
# yapf:disable
log_config = dict(
    interval=1,
    hooks=[
        dict(type='TextLoggerHook'),
        # dict(type='TensorboardLoggerHook')
    ])
# yapf:enable
dist_params = dict(backend='nccl')
log_level = 'INFO'
load_from = None
resume_from = None
workflow = [('train', 1)]
1
2
3
4
5
6
7
8
9
10
11
12
13
14

训练

通过运行如下命令，即可开启Mask R-CNN的训练。

CUDA_VISIBLE_DEVICES=4,5,6,7 \
bash tools/dist_train.sh configs/aaaa/mask_rcnn_r50_fpn_custom.py 4
1
2

测试

通过运行test.py文件，来开启单GPU的测试，命令如下。

python tools/test.py /path/to/config_file /path/to/checkpoint_file --eval bbox segm
1

常见问题

Q1：oserror: [errno 39] directory not empty "eval_hook"

通过注释mmdet/core/evaluation/eval_hooks.py文件中的tmpdir内容，具体操作是将multi_gpu_test函数中的tmpdir设置为None。

results = multi_gpu_test(
    runner.model,
    self.dataloader,
    # tmpdir=tmpdir,
    tmpdir=None,
    gpu_collect=self.gpu_collect)
1
2
3
4
5
6

参考链接

【实例分割（一）】Detectron2 数据集制作并注册数据集训练 - 古月居
【实例分割（二）】Mask2Former 数据集制作和训练 - 古月居
【深度学习】YOLOv5实例分割 数据集制作、模型训练以及TensorRT部署
利用labelme制作实例分割数据集_labelme实例分割_Jiazhou_garland的博客-CSDN博客