MMDetection2（二）：自定义数据集

目录

一、支持新的数据格式

1、将新的数据格式构造为已存在的格式

2、将新的数据格式构造为中间格式

3、一个自定义数据集的例子

二、通过dataset wrappers自定义数据集

三、更改数据集的类

总结

---

一、支持新的数据格式

为了支持新的数据格式，你能够将它们转化为已经存在的VOC和COCO格式，或者直接将它们转化为中间格式。你可以选择在线下将它们转化或者在线上转换。在MMDetection中，我们推荐将数据集在线下转化为COCO格式，因此你仅仅需要更改配置文件中的data annotation的路径和classes就好了。

1、将新的数据格式构造为已存在的格式

最简单的方式是将你的数据集转化为已存在的数据集格式（例如COCO或者PASCAL VOC）

COCO格式数据集的json标注文件含有以下几个必要的关键字：

'images': [
    {
        'file_name': 'COCO_val2014_000000001268.jpg',
        'height': 427,
        'width': 640,
        'id': 1268
    },
    ...
],

'annotations': [
    {
        'segmentation': [[192.81,
            247.09,
            ...
            219.03,
            249.06]],  # 如果你有mask labels的话
        'area': 1035.749,
        'iscrowd': 0,
        'image_id': 1268,
        'bbox': [192.81, 224.8, 74.73, 33.43],
        'category_id': 16,
        'id': 42986
    },
    ...
],

'categories': [
    {'id': 0, 'name': 'car'},
 ]

在json文件中有三个必要的key值：

1. images:包含有图片信息的list，例如file_name,height,width以及id。
2. annotations:包含有实例标注信息的list。
3. categories:包含有的categories名和对应ID的list。

在进行了数据预处理之后，使用者需要更改配置文件从而来使用数据集。接下来的例子是使用一个有5个类别的自定义数据集，假设它也是COCO格式。

In configs/my_custom_config.py:

...
# 数据集的设置
dataset_type = 'CocoDataset'
classes = ('a', 'b', 'c', 'd', 'e')
...
data = dict(
    samples_per_gpu=2,
    workers_per_gpu=2,
    train=dict(
        type=dataset_type,   # 数据集的类型
        classes=classes,
        ann_file='path/to/your/train/data',   # 训练集标注文件的路径  
        ...),
    val=dict(
        type=dataset_type,  # 数据集的类型
        classes=classes,
        ann_file='path/to/your/val/data',  # 验证集标注文件的路径
        ...),
    test=dict(
        type=dataset_type,  # 数据集的类型
        classes=classes,
        ann_file='path/to/your/test/data',  # 测试集标注文件的路径
        ...))
...

2、将新的数据格式构造为中间格式

略

这个地方我没怎么用过，就跳过了~~~，有兴趣的话可以去看原文档orz

3、一个自定义数据集的例子

假定数据集的标签是一种新的文本文件格式。Bounding boxes标签信息通过以下形式存储在annotation.txt文件中:

#
000001.jpg
1280 720
2
10 20 40 60 1
20 40 50 60 2
#
000002.jpg
1280 720
3
50 20 40 60 2
20 40 30 45 2
30 40 50 60 3

我们创建一个新的数据集文件mmdet/datasets/my_dataset.py来导入数据。

import mmcv
import numpy as np

from .builder import DATASETS
from .custom import CustomDataset


@DATASETS.register_module()
class MyDataset(CustomDataset):

    CLASSES = ('person', 'bicycle', 'car', 'motorcycle')

    def load_annotations(self, ann_file):
        ann_list = mmcv.list_from_file(ann_file)

        data_infos = []
        for i, ann_line in enumerate(ann_list):
            if ann_line != '#':
                continue

            img_shape = ann_list[i + 2].split(' ')
            width = int(img_shape[0])
            height = int(img_shape[1])
            bbox_number = int(ann_list[i + 3])

            anns = ann_line.split(' ')
            bboxes = []
            labels = []
            for anns in ann_list[i + 4:i + 4 + bbox_number]:
                bboxes.append([float(ann) for ann in anns[:4]])
                labels.append(int(anns[4]))

            data_infos.append(
                dict(
                    filename=ann_list[i + 1],
                    width=width,
                    height=height,
                    ann=dict(
                        bboxes=np.array(bboxes).astype(np.float32),
                        labels=np.array(labels).astype(np.int64))
                ))

        return data_infos

    def get_ann_info(self, idx):
        return self.data_infos[idx]['ann']

然后在配置文件中去使用MyDataset，你可以按照下面的这种样式更改配置文件：

dataset_A_train = dict(
    type='MyDataset',
    ann_file = 'image_list.txt',
    pipeline=train_pipeline
)

二、通过dataset wrappers自定义数据集

略

这一部分因为我自己也没太用过，所以就跳过orz~~~，有兴趣的可以去看原文档。

三、更改数据集的类

使用现有的数据集类型，我们可以更改它们的类名来训练标注数据的子集。举一个栗子，如果你仅仅是只想训练当前数据集的三个类，那么你只需要更改配置文件中数据集的类别参数就好了。

数据集导入的时候会自动过滤掉其他几个类别的标注。

# 如果你只需要对person、bicycle、car进行检测
classes = ('person', 'bicycle', 'car')
# 在配置文件中更改classes参数的值
data = dict(
    train=dict(classes=classes),
    val=dict(classes=classes),
    test=dict(classes=classes))

 MMDetection V2.0开始支持从一个文件中读取classes，假定classes.txt文件中包含有以下几个类名:

person
bicycle
car

使用者可以设置classes的值设为文件的路径，数据集将会自动的导入并转化它为一个列表。

classes = 'path/to/classes.txt'
data = dict(
    train=dict(classes=classes),
    val=dict(classes=classes),
    test=dict(classes=classes))

总结

这一节主要是介绍如何自定义数据集，关键是看一下自定义数据集的那个例子，需要自己写一个继承CustomDataset的类，然后在配置文件中使用这个类。因为中间格式和dataset wrappers我没怎么用过，所以我也就没怎么看这一块，有兴趣的可以自己去看一下原文档。