妙啊！MMDetection 训练自定义数据集

深度学习发展到现在已经有很多优秀的模型，同时很多大公司也会在内部开发自己的框架，快速的实现业务，从而产生实际价值。

如下面的招聘要求一样，市场需要这些能熟练使用现有工具快速实现，MMDetection 是一个非常好的选择。

为了便于记录和理解内容，这里将整篇文章的主要内容，绘制成思维导图的形式：
接下来开始实际使用！如果对你有所帮助，请给我意见三连。

在本文中，你将知道如何使用定制的数据集推断、测试和训练预定义的模型。我们以ballon数据集为例来描述整个过程。

气球数据集：https://github.com/matterport/Mask_RCNN/tree/master/samples/balloon

https://github.com/matterport/Mask_RCNN/releases

1、准备自定义数据集

官方教程：https://mmdetection.readthedocs.io/en/latest/2_new_data_model.html

有三种方法在MMDetection中支持新的数据集:

  • 将数据集重新组织为COCO格式。

  • 将数据集重新组织为中间格式。

  • 实现一个新的数据集。

官方建议使用前两种方法，这两种方法通常比第三种方法简单。

在本文中，我们给出了一个将数据转换为COCO格式的示例。

注意：MMDetection目前只支持评估COCO格式数据集的mask AP。因此，例如实例分割任务，用户应该将数据转换为coco格式。

COCO 标注格式

以下是实例分割所需的COCO格式所需的关键，完整的细节请参考这里。

https://cocodataset.org/#format-data

{
'images': [image],
'annotations': [annotation],
'categories': [category]
}

image = {
'id': int,
'width': int,
'height': int,
'file_name': str,
}

annotation = {
'id': int,
'image_id': int,
'category_id': int,
'segmentation': RLE or [polygon],
'area': float,
'bbox': [x,y,width,height],
'iscrowd': 0 or 1,
}

categories = [{
'id': int,
'name': str,
'supercategory': str,
}]

假设我们使用ballon数据集。下载数据之后，我们需要实现一个函数来将注释格式转换为COCO格式。然后我们可以使用实现的COCODataset加载数据，并执行训练和评估。

如果你看一下数据集，你会发现数据集的格式如下:

{'base64_img_data': '', 'file_attributes': {}, 'filename': '34020010494_e5cb88e1c4_k.jpg', 'fileref': '', 'regions': {'0': {'region_attributes': {},   'shape_attributes': {'all_points_x': [1020,     1000,     994,     1003,     1023,     1050,     1089,     1134,     1190,     1265,     1321,     1361,     1403,     1428,     1442,     1445,     1441,     1427,     1400,     1361,     1316,     1269,     1228,     1198,     1207,     1210,     1190,     1177,     1172,     1174,     1170,     1153,     1127,     1104,     1061,     1032,     1020],    'all_points_y': [963,     899,     841,     787,     738,     700,     663,     638,     621,     619,     643,     672,     720,     765,     800,     860,     896,     942,     990,     1035,     1079,     1112,     1129,     1134,     1144,     1153,     1166,     1166,     1150,     1136,     1129,     1122,     1112,     1084,     1037,     989,     963],    'name': 'polygon'}}}, 'size': 1115004}

annotation 是一个JSON文件，其中每个 key 都表示图像的所有注释。将ballon数据集转换为coco格式的代码如下所示。

import os.path as ospimport mmcv
def convert_balloon_to_coco(ann_file, out_file, image_prefix): data_infos = mmcv.load(ann_file)
 annotations = [] images = [] obj_count = 0
 for idx, v in enumerate(mmcv.track_iter_progress(data_infos.values())): filename = v['filename'] img_path = osp.join(image_prefix, filename) height, width = mmcv.imread(img_path).shape[:2]
 images.append(dict( id = idx, file_name = filename, height=height, width = width))
 bboxes = [] labels = [] masks = [] for _, obj in v['regions'].items(): assert not obj['region_attributes'] obj = obj['shape_attributes'] px = obj['all_points_x'] py = obj['all_points_y'] poly = [(x+0.5, y+0.5) for x,y in zip(px,py)] poly = [p for x in poly for p in x]
 x_min, y_min, x_max, y_max = ( min(px), min(py), max(px),max(py)) data_anno = dict( image_id = idx, id = obj_count, category_id = 0, bbox = [x_min, y_min, x_max-x_min, y_max-y_min], area = (x_max - x_min)*(y_max - y_min), segmentation = [poly], iscrowd =0)
 annotations.append(data_anno) obj_count += 1 coco_format_json = dict( images = images, annotations = annotations, categories=[{'id':0, 'name':'balloon'}] ) mmcv.dump(coco_format_json, out_file)
# 对验证集数据进行处理是，将下面路径中的train 替换成val 即可# 注意数据集 balloon 的路径自行调整ann_file = './balloon/train/via_region_data.json'out_file = './balloon/train/annotation_coco.json'image_prefix = './balloon/train'convert_balloon_to_coco(ann_file, out_file, image_prefix)

注释：

# 可以加载 json, yaml, pkl 文件import mmcvmmcv.load('test.json')

# 刷新位置的进度条方式mmcv.track_iter_progress(tasks)

参考资料：https://zhuanlan.zhihu.com/p/126725557

https://mmcv.readthedocs.io/en/stable/

通过上面的函数，用户可以成功地将标注文件转换成json格式，然后我们可以使用CocoDataset对模型进行训练和评估。

2、config文件配置

第二步是准备一个 config，这样数据集就可以成功加载。假设我们想使用带有FPN的Mask R-CNN，在balloon数据集上训练检测器的配置如下。假设配置在configs/balloon/目录下，命名为mask_rcnn_r50_caffe_fpn_mstrain-poly_1x_balloon.py。配置如下所示。

# The new config inherits a base config to highlight the necessary modification_base_ = '../mask_rcnn/mask_rcnn_r50_caffe_fpn_mstrain-poly_1x_coco.py'
# We also need to change the num_classes in head to match the dataset's annotationmodel = dict( roi_head=dict( bbox_head=dict(num_classes=1), mask_head=dict(num_classes=1)))
# Modify dataset related settingsdataset_type = 'COCODataset'classes = ('balloon',)data = dict( train=dict( img_prefix='balloon/train/', classes=classes, ann_file='balloon/train/annotation_coco.json'), val=dict( img_prefix='balloon/val/', classes=classes, ann_file='balloon/val/annotation_coco.json'), test=dict( img_prefix='balloon/val/', classes=classes, ann_file='balloon/val/annotation_coco.json'))
# We can use the pre-trained Mask RCNN model to obtain higher performanceload_from = 'checkpoints/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco_bbox_mAP-0.408__segm_mAP-0.37_20200504_163245-42aa3d00.pth'

注意：

这里的_base_ 要修改成

_base_ = '../mask_rcnn/mask_rcnn_r50_caffe_fpn_mstrain-poly_1x_coco.py'官方提供的路径有一点问题

3、自定义数据集上训练、测试、推理模型

训练一个新模型

使用新的config 训练一个模型，直接运行下面的代码即可：

python tools/train.py configs/balloon/mask_rcnn_r50_caffe_fpn_mstrain-poly_1x_balloon.py

如果报错

 raise IOError(f'{filename} is not a checkpoint file')OSError: checkpoints/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco_bbox_mAP-0.408__segm_mAP-0.37_20200504_163245-42aa3d00.pth is not a checkpoint file

建议去官方提供的预训练模型下载地址去下载，并放置在checkpoints 文件夹下

https://mmdetection.readthedocs.io/en/latest/model_zoo.html

直接下载：http://download.openmmlab.com/mmdetection/v2.0/mask_rcnn/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco_bbox_mAP-0.408__segm_mAP-0.37_20200504_163245-42aa3d00.pth

注意：

大概需要 9 G 的现存才能跑的起来。。。

测试并推理

测试训练好的模型，直接运行：

python tools/test.py configs/balloon/mask_rcnn_r50_caffe_fpn_mstrain-poly_1x_balloon.py work_dirs/mask_rcnn_r50_caffe_fpn_mstrain-poly_1x_balloon.py/latest.pth --eval bbox segm
赞 (0)

相关推荐