Detectron2 trains its own dataset and converts it to coco format

Reference resources

About coco The format of
https://detectron2.readthedocs.io/en/latest/tutorials/datasets.html#register-a-dataset

Register and train your data set https://blog.csdn.net/qq_29750461/article/details/106761382

https://cloud.tencent.com/developer/article/1960793
coco api https://blog.csdn.net/qq_41709370/article/details/108471072

Note that if there is an error saying that your data is not registered , Need to be in main Add the function to it plain_register_dataset() Register your own dataset

Data to coco Format

Some of the pit ：

• bbox about coco Format [x,y,width,height] x,y Is the coordinate of the upper left corner

• file_name Is the absolute path of the picture To write xxx.jpg

• There are some coco The file says that it is not a required field, but if there is no field, it will report keyerror
  anntations In such as area、iscrowd Default values can be set area It can be bbox Of w*h

• Read the picture shape[0] yes height shape[1] yes witdh

height = img.shape[0]
width = img.shape[1]

Register data sets

#!/usr/bin/env python
# Copyright (c) Facebook, Inc. and its affiliates.
""" A main training script. This scripts reads a given config file and runs the training or evaluation. It is an entry point that is made to train standard models in detectron2. In order to let one script support training of many models, this script contains logic that are specific to these built-in models and therefore may not be suitable for your own project. For example, your research project perhaps only needs a single "evaluator". Therefore, we recommend you to use detectron2 as an library and take this file as an example of how to use the library. You may want to write your own script with your datasets and other customizations. """

import logging
import os
from collections import OrderedDict

import cv2
from detectron2.utils.visualizer import Visualizer

import detectron2.utils.comm as comm
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.config import get_cfg
from detectron2.data import MetadataCatalog
from detectron2.engine import DefaultTrainer, default_argument_parser, default_setup, hooks, launch
from detectron2.evaluation import (
    CityscapesInstanceEvaluator,
    CityscapesSemSegEvaluator,
    COCOEvaluator,
    COCOPanopticEvaluator,
    DatasetEvaluators,
    LVISEvaluator,
    PascalVOCDetectionEvaluator,
    SemSegEvaluator,
    verify_results,
)
from detectron2.modeling import GeneralizedRCNNWithTTA

from detectron2.data import DatasetCatalog, MetadataCatalog
from detectron2.data.datasets.coco import load_coco_json
import pycocotools
# Claim categories , Try to keep 
CLASS_NAMES =["background", "military"]
#  Dataset path 
DATASET_ROOT = '/home/szr/new_dete/detectron2/datasets'
ANN_ROOT = os.path.join(DATASET_ROOT, 'COCOformat')

TRAIN_PATH = os.path.join(DATASET_ROOT, 'JPEGImages')
VAL_PATH = os.path.join(DATASET_ROOT, 'JPEGImages')

TRAIN_JSON = os.path.join(ANN_ROOT, 'train.json')
VAL_JSON = os.path.join(ANN_ROOT, 'val.json')
# VAL_JSON = os.path.join(ANN_ROOT, 'test.json')

#  Declare a subset of the dataset 
PREDEFINED_SPLITS_DATASET = {
    
    "coco_my_train": (TRAIN_PATH, TRAIN_JSON),
    "coco_my_val": (VAL_PATH, VAL_JSON),
}
#=========== There are two ways to register data sets , The second one I use directly plain_register_dataset The way   It can also be used. register_dataset In the form of ==================
# Register data sets （ This step is to register the custom dataset into Detectron2）

#  Register datasets and metadata 
def plain_register_dataset():
    # Training set 
    DatasetCatalog.register("coco_my_train", lambda: load_coco_json(TRAIN_JSON, TRAIN_PATH))
    MetadataCatalog.get("coco_my_train").set(thing_classes=CLASS_NAMES,  #  You can choose to turn on , But it cannot display Chinese , Here we need to pay attention to , In Chinese, it's better to close 
                                                    evaluator_type='coco', #  Specify the evaluation method 
                                                    json_file=TRAIN_JSON,
                                                    image_root=TRAIN_PATH)

    #DatasetCatalog.register("coco_my_val", lambda: load_coco_json(VAL_JSON, VAL_PATH, "coco_2017_val"))
    # verification / Test set 
    DatasetCatalog.register("coco_my_val", lambda: load_coco_json(VAL_JSON, VAL_PATH))
    MetadataCatalog.get("coco_my_val").set(thing_classes=CLASS_NAMES, #  You can choose to turn on , But it cannot display Chinese , Here we need to pay attention to , In Chinese, it's better to close 
                                                evaluator_type='coco', #  Specify the evaluation method 
                                                json_file=VAL_JSON,
                                                image_root=VAL_PATH)


#  View dataset annotations , Visually check whether the dataset annotation is correct ,
# You can also write your own script to judge , In fact, it is to judge whether the callout box exceeds the image boundary 
# You can choose to use this method 
def checkout_dataset_annotation(name="coco_my_val"):
    #dataset_dicts = load_coco_json(TRAIN_JSON, TRAIN_PATH, name)
    dataset_dicts = load_coco_json(TRAIN_JSON, TRAIN_PATH)
    print(len(dataset_dicts))
    for i, d in enumerate(dataset_dicts,0):
        #print(d)
        img = cv2.imread(d["file_name"])
        visualizer = Visualizer(img[:, :, ::-1], metadata=MetadataCatalog.get(name), scale=1.5)
        vis = visualizer.draw_dataset_dict(d)
        #cv2.imshow('show', vis.get_image()[:, :, ::-1])
        cv2.imwrite('out/'+str(i) + '.jpg',vis.get_image()[:, :, ::-1])
        #cv2.waitKey(0)
        if i == 200:
            break


def build_evaluator(cfg, dataset_name, output_folder=None):
    """ Create evaluator(s) for a given dataset. This uses the special metadata "evaluator_type" associated with each builtin dataset. For your own dataset, you can simply create an evaluator manually in your script and do not have to worry about the hacky if-else logic here. """
    if output_folder is None:
        output_folder = os.path.join(cfg.OUTPUT_DIR, "inference")
    evaluator_list = []
    evaluator_type = MetadataCatalog.get(dataset_name).evaluator_type
    if evaluator_type in ["sem_seg", "coco_panoptic_seg"]:
        evaluator_list.append(
            SemSegEvaluator(
                dataset_name,
                distributed=True,
                output_dir=output_folder,
            )
        )
    if evaluator_type in ["coco", "coco_panoptic_seg"]:
        evaluator_list.append(COCOEvaluator(dataset_name, output_dir=output_folder))
    if evaluator_type == "coco_panoptic_seg":
        evaluator_list.append(COCOPanopticEvaluator(dataset_name, output_folder))
    if evaluator_type == "cityscapes_instance":
        return CityscapesInstanceEvaluator(dataset_name)
    if evaluator_type == "cityscapes_sem_seg":
        return CityscapesSemSegEvaluator(dataset_name)
    elif evaluator_type == "pascal_voc":
        return PascalVOCDetectionEvaluator(dataset_name)
    elif evaluator_type == "lvis":
        return LVISEvaluator(dataset_name, output_dir=output_folder)
    if len(evaluator_list) == 0:
        raise NotImplementedError(
            "no Evaluator for the dataset {} with the type {}".format(dataset_name, evaluator_type)
        )
    elif len(evaluator_list) == 1:
        return evaluator_list[0]
    return DatasetEvaluators(evaluator_list)


class Trainer(DefaultTrainer):
    """ We use the "DefaultTrainer" which contains pre-defined default logic for standard training workflow. They may not work for you, especially if you are working on a new research project. In that case you can write your own training loop. You can use "tools/plain_train_net.py" as an example. """

    @classmethod
    def build_evaluator(cls, cfg, dataset_name, output_folder=None):
        return build_evaluator(cfg, dataset_name, output_folder)

    @classmethod
    def test_with_TTA(cls, cfg, model):
        logger = logging.getLogger("detectron2.trainer")
        # In the end of training, run an evaluation with TTA
        # Only support some R-CNN models.
        logger.info("Running inference with test-time augmentation ...")
        model = GeneralizedRCNNWithTTA(cfg, model)
        evaluators = [
            cls.build_evaluator(
                cfg, name, output_folder=os.path.join(cfg.OUTPUT_DIR, "inference_TTA")
            )
            for name in cfg.DATASETS.TEST
        ]
        res = cls.test(cfg, model, evaluators)
        res = OrderedDict({
    k + "_TTA": v for k, v in res.items()})
        return res


# def setup(args):
# """
# Create configs and perform basic setups.
# """
# # cfg = get_cfg()
# # cfg.merge_from_file(args.config_file)
# # cfg.merge_from_list(args.opts)
# # cfg.freeze()
# # default_setup(cfg, args)
# # return cfg

def setup(args):
    """ Create configs and perform basic setups. """
    cfg = get_cfg()
    args.config_file = "/home/szr/new_dete/detectron2/configs/COCO-Detection/retinanet_R_50_FPN_3x.yaml"
    cfg.merge_from_file(args.config_file)  #  from config file  Override configuration 
    cfg.merge_from_list(args.opts)  #  from CLI Parameters   Override configuration 

    #  Change configuration parameters 
    cfg.DATASETS.TRAIN = ("coco_my_train",)  #  Training dataset name 
    cfg.DATASETS.TEST = ("coco_my_val",)
    cfg.DATALOADER.NUM_WORKERS = 4  #  Single thread 

    cfg.INPUT.CROP.ENABLED = True
    cfg.INPUT.MAX_SIZE_TRAIN = 640  #  Maximum size of training picture input 
    cfg.INPUT.MAX_SIZE_TEST = 640  #  Maximum size of test data input 
    cfg.INPUT.MIN_SIZE_TRAIN = (512, 768)  #  Minimum size of training image input , It can be set as multi-scale training 
    cfg.INPUT.MIN_SIZE_TEST = 640
    # cfg.INPUT.MIN_SIZE_TRAIN_SAMPLING, There are two configurations , Respectively  choice  And  range ：
    # range  Let the short edge of the image from  512-768 Random selection 
    # choice ：  Convert the input image to the specified , A limited number of image sizes for training , That is, the short side can only be  512 perhaps 768
    cfg.INPUT.MIN_SIZE_TRAIN_SAMPLING = 'range'
    #  Be sure to read the notes ！！！！！！！！
    cfg.MODEL.RETINANET.NUM_CLASSES = 2  #  Number of categories +1（ Because there is background, It's your  cate id  from  1  Start , If your dataset Json Subscript from  0  Start , Change this to your corresponding category , No more background classes ！！！！！）
    # cfg.MODEL.WEIGHTS="/home/yourstorePath/.pth"
    # This address can be downloaded by itself or it can be downloaded online by itself 
    cfg.MODEL.WEIGHTS = "detectron2://COCO-Detection/retinanet_R_50_FPN_3x/190397829/model_final_5bd44e.pkl"  #  Pre training model weight 
    cfg.SOLVER.IMS_PER_BATCH = 4  # batch_size=2; iters_in_one_epoch = dataset_imgs/batch_size

    #  According to the total number of training data and batch_size, Work out each epoch Number of iterations required 
    # 9000 The total number of data for your workout , Customizable 
    ITERS_IN_ONE_EPOCH = int(1120 / cfg.SOLVER.IMS_PER_BATCH)

    #  Specify the maximum number of iterations 
    cfg.SOLVER.MAX_ITER = (ITERS_IN_ONE_EPOCH * 12) - 1  # 12 epochs,
    #  Initial learning rate 
    cfg.SOLVER.BASE_LR = 0.002
    #  Optimizer kinetic energy 
    cfg.SOLVER.MOMENTUM = 0.9
    #  Weight falloff 
    cfg.SOLVER.WEIGHT_DECAY = 0.0001
    cfg.SOLVER.WEIGHT_DECAY_NORM = 0.0
    #  Learning rate attenuation multiple 
    cfg.SOLVER.GAMMA = 0.1
    #  Iterate to a specified number of times , Learning rate is attenuated 
    cfg.SOLVER.STEPS = (7000,)
    #  Before training , Will do a warm-up exercise , The learning rate slowly increases the initial learning rate 
    cfg.SOLVER.WARMUP_FACTOR = 1.0 / 1000
    #  Warm up iterations 
    cfg.SOLVER.WARMUP_ITERS = 1000

    cfg.SOLVER.WARMUP_METHOD = "linear"
    #  Save the named data of the model file 1
    cfg.SOLVER.CHECKPOINT_PERIOD = ITERS_IN_ONE_EPOCH - 1

    #  Iterate to a specified number of times , Conduct an assessment 
    cfg.TEST.EVAL_PERIOD = ITERS_IN_ONE_EPOCH
    # cfg.TEST.EVAL_PERIOD = 100

    # cfg.merge_from_file(args.config_file)
    # cfg.merge_from_list(args.opts)
    cfg.freeze()
    default_setup(cfg, args)
    return cfg


def main(args):
    cfg = setup(args)

    if args.eval_only:
        model = Trainer.build_model(cfg)
        DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
            cfg.MODEL.WEIGHTS, resume=args.resume
        )
        res = Trainer.test(cfg, model)
        if cfg.TEST.AUG.ENABLED:
            res.update(Trainer.test_with_TTA(cfg, model))
        if comm.is_main_process():
            verify_results(cfg, res)
        return res

    """ If you'd like to do anything fancier than the standard training logic, consider writing your own training loop (see plain_train_net.py) or subclassing the trainer. """
    trainer = Trainer(cfg)
    trainer.resume_or_load(resume=args.resume)
    if cfg.TEST.AUG.ENABLED:
        trainer.register_hooks(
            [hooks.EvalHook(0, lambda: trainer.test_with_TTA(cfg, trainer.model))]
        )
    return trainer.train()


if __name__ == "__main__":
    args = default_argument_parser().parse_args()
    plain_register_dataset()
    print("Command Line Args:", args)
    launch(
        main,
        args.num_gpus,
        num_machines=args.num_machines,
        machine_rank=args.machine_rank,
        dist_url=args.dist_url,
        args=(args,),
    )