ubuntu系统 detectron2训练自己的数据库

@ detectron2训练自己的数据库

注册自己的数据集

detectron2官网地址 https://github.com/facebookresearch/detectron2。安装及环境配置参考官方文档或者其他博客，此处不在赘述，主要描述训练自己数据集的过程。本文使用的数据集为visdrone2018，在训练数据集之前需要将数据集转换为coco格式的json数据集。

方法1

from detectron2.data.datasets import register_coco_instances

register_coco_instances('self_coco_train', {
   
   }, 
                        'instances_train2017.json', #json文件路径
                       'train2017'#图片路径)
register_coco_instances('self_coco_val', {
   
   }, 
                        'instances_val2017.json', #json文件路径
                       'val2017'#图片路径)

方法2

from detectron2.data import DatasetCatalog, MetadataCatalog
from detectron2.data.datasets.coco import load_coco_json

DatasetCatalog.register("self_coco_train", lambda: load_coco_json(TRAIN_JSON, TRAIN_PATH, "self_coco_train"))
MetadataCatalog.get("self_coco_train").set(thing_classes=["pos", "neg"],
                                                    json_file=TRAIN_JSON,
                                                    image_root=TRAIN_PATH)
DatasetCatalog.register("self_coco_val", lambda: load_coco_json(VAL_JSON, VAL_PATH, "self_coco_val"))
MetadataCatalog.get("self_coco_val").set(thing_classes=["pos", "neg"],
                                                json_file=VAL_JSON,
                                                image_root=VAL_PATH)

获取数据的元数据

coco_val_metadata = MetadataCatalog.get("self_coco_val")
dataset_dicts = DatasetCatalog.get("self_coco_val")
coco_val_metadata

数据可视化，确定数据载入无问题

import random

for d in random.sample(dataset_dicts, 3):
    img = cv2.imread(d["file_name"])
    visualizer = Visualizer(img[:, :, ::-1], metadata=coco_val_metadata, scale=0.5)
    vis = visualizer.draw_dataset_dict(d)
    cv2.imshow(vis.get_image()[:, :, ::-1])
    cv2.waitKey(0)

完整代码

根据官方的tools/train_net.py 修改

#-*- coding:utf-8 -*-
import os
import cv2
import logging
from collections import OrderedDict

import detectron2.utils.comm as comm
from detectron2.utils.visualizer import Visualizer
from detectron2.checkp