使用Detectron2框架训练自己的数据集

一. 引言

最近在使用detectron2框架学习一些开源算法，一上来，数据加载把我搞的有点懵，再次记录一下，如何修改代码，训练自己的数据集。

二. detectron2 框架安装

git clone https://github.com/facebookresearch/detectron2.git
python -m pip install -e detectron2

三. 数据集注册

本次我的数据是coco数据格式，所以针对coco数据修改代码，其他数据形式，也是类似。
关于数据集的注册和流程可以看看源码，以下4个函数

data/datasets/builtin.py/register_all_coco（）
data/datasets/builtin_meta.py/_get_builtin_metadata（）
data/datasets/builtin_meta.py/_get_coco_instances_meta() 
data/datasets/register_coco.py/register_coco_instances()

对于数据存储路径格式可以参考 detectron2/datasets/README.md 本次是coco数据格式

coco/
  annotations/
    instances_{train,val}2017.json
    person_keypoints_{train,val}2017.json
  {train,val}2017/
    # image files that are mentioned in the corresponding json

注册数据集（本次实验数据集，数据形式是coco格式存储，其余格式需要按照其他格式存储）
   vim ～/detectron2/detectron2/data/datasets/builtin.py #注册数据集
   PREDEFINED_SPLITS_COCO["coco"] = {
    "coco_2014_train": ("coco/train2014", "coco/annotations/instances_train2014.json"),
    "coco_2014_val": ("coco/val2014", "coco/annotations/instances_val2014.json"),
    "coco_2014_minival": ("coco/val2014", "coco/annotations/instances_minival2014.json"),
    "coco_2014_valminusminival": (
        "coco/val2014",
        "coco/annotations/instances_valminusminival2014.json",
    ),
    "coco_2017_train": ("coco/train2017", "coco/annotations/instances_train2017.json"),
    "coco_2017_val": ("coco/val2017", "coco/annotations/instances_val2017.json"),
    "coco_2017_test": ("coco/test2017", "coco/annotations/image_info_test2017.json"),
    "coco_2017_test-dev": ("coco/test2017", "coco/annotations/image_info_test-dev2017.json"),
    "coco_2017_val_100": ("coco/val2017", "coco/annotations/instances_val2017_100.json"),
    
    #我新注册的数据集，注意路径别写错了.
    "coco_zj_train": ('coco/annotations/images', "coco/annotations/zj_train.json"),
    "coco_zj_val": ('coco/annotations/images', 'coco/annotations/zj_val.json'),
    # "coco_zj_train" : 数据集名称
    # 'coco/annotations/images' ： 图片存放路径 /coco/annotations/images #datasets前面不用说明
    # 'coco/annotations/zj_train.json' : 标注信息json路径 /coco/annotations/zj_train.json #datasets前面不用说明
}

 vim ～/detectron2/detectron2/data/datasets/builtin_meta.py #给类别定义颜色

 COCO_CATEGORIES_zj = [
    {"color": [220, 20, 60], "isthing": 1, "id": 1, "name": "person"},
    {"color": [119, 11, 32], "isthing": 1, "id": 2, "name": "bicycle"},
    {"color": [0, 0, 142], "isthing": 1, "id": 3, "name": "car"},        
]
   def _get_coco_instances_meta_zj():
    thing_ids = [k["id"] for k in COCO_CATEGORIES_zj if k["isthing"] == 1]
    thing_colors = [k["color"] for k in COCO_CATEGORIES_zj if k["isthing"] == 1]
    assert len(thing_ids) == 3, len(thing_ids)
    # Mapping from the incontiguous COCO category id to an id in [0, 79]
    thing_dataset_id_to_contiguous_id = {k: i for i, k in enumerate(thing_ids)}
    thing_classes = [k["name"] for k in COCO_CATEGORIES_zj if k["isthing"] == 1]
    '''
        _get_coco_instances_meta thing_classes: ['person', 'bicycle', 'car']
    '''    
    ret = {
        "thing_dataset_id_to_contiguous_id": thing_dataset_id_to_contiguous_id,
        "thing_classes": thing_classes,
        "thing_colors": thing_colors,
    }
    '''
        ret = {'thing_dataset_id_to_contiguous_id': {1: 0, 2: 1, 3: 2}, 'thing_classes': ['person', 'bicycle', 'car'], 'thing_colors': [[119, 11, 32], [0, 0, 142], [220, 20, 60]]}
    ''' 
    return ret

为数据集注册元数

def _get_builtin_metadata(dataset_name):
    if dataset_name == "coco":
        return _get_coco_instances_meta_zj() #修改成刚才新命名的函数名
    if dataset_name == "coco_panoptic_separated":
        return _get_coco_panoptic_separated_meta()
    elif dataset_name == "coco_panoptic_standard":
        meta = {}
        # The following metadata maps contiguous id from [0, #thing categories +
        # #stuff categories) to their names and colors. We have to replica of the
        # same name and color under "thing_*" and "stuff_*" because the current
        # visualization function in D2 handles thing and class classes differently
        # due to some heuristic used in Panoptic FPN. We keep the same naming to
        # enable reusing existing visualization functions.
        thing_classes = [k["name"] for k in COCO_CATEGORIES]
        thing_colors = [k["color"] for k in COCO_CATEGORIES]
        stuff_classes = [k["name"] for k in COCO_CATEGORIES]
        stuff_colors = [k["color"] for k in COCO_CATEGORIES]

        meta["thing_classes"] = thing_classes
        meta["thing_colors"] = thing_colors
        meta["stuff_classes"] = stuff_classes
        meta["stuff_colors"] = stuff_colors

        # Convert category id for training:
        #   category id: like semantic segmentation, it is the class id for each
        #   pixel. Since there are some classes not used in evaluation, the category
        #   id is not always contiguous and thus we have two set of category ids:
        #       - original category id: category id in the original dataset, mainly
        #           used for evaluation.
        #       - contiguous category id: [0, #classes), in order to train the linear
        #           softmax classifier.
        thing_dataset_id_to_contiguous_id = {}
        stuff_dataset_id_to_contiguous_id = {}

        for i, cat in enumerate(COCO_CATEGORIES):
            if cat["isthing"]:
                thing_dataset_id_to_contiguous_id[cat["id"]] = i
            else:
                stuff_dataset_id_to_contiguous_id[cat["id"]] = i

        meta["thing_dataset_id_to_contiguous_id"] = thing_dataset_id_to_contiguous_id
        meta["stuff_dataset_id_to_contiguous_id"] = stuff_dataset_id_to_contiguous_id

        return meta
    elif dataset_name == "coco_person":
        return {
            "thing_classes": ["person"],
            "keypoint_names": COCO_PERSON_KEYPOINT_NAMES,
            "keypoint_flip_map": COCO_PERSON_KEYPOINT_FLIP_MAP,
            "keypoint_connection_rules": KEYPOINT_CONNECTION_RULES,
        }
    elif dataset_name == "cityscapes":
        # fmt: off
        CITYSCAPES_THING_CLASSES = [
            "person", "rider", "car", "truck",
            "bus", "train", "motorcycle", "bicycle",
        ]
        CITYSCAPES_STUFF_CLASSES = [
            "road", "sidewalk", "building", "wall", "fence", "pole", "traffic light",
            "traffic sign", "vegetation", "terrain", "sky", "person", "rider", "car",
            "truck", "bus", "train", "motorcycle", "bicycle",
        ]
        # fmt: on
        return {
            "thing_classes": CITYSCAPES_THING_CLASSES,
            "stuff_classes": CITYSCAPES_STUFF_CLASSES,
        }
    raise KeyError("No built-in metadata for dataset {}".format(dataset_name))

其余配置文件yaml,参考官方案例，都是详细的说明

开始训练
python detectron2/tools/train_net.py

可以看到以下输出日志