Several pitfalls encountered when running bmask-rcnn

Source code address

hustvl/BMaskR-CNN: Boundary-preserving Mask R-CNN (ECCV 2020) (github.com)

1. Problem encountered 1: Installing detectron2

Please be sure to follow the official tutorial for installation, especially when using windows + anaconda. It is recommended to install from source code.

Never pip

Official installation tutorial:

Use Custom Datasets – detectron2 0.6 documentation

2. Problem 2: Missing file bmask_rcnn_R_50_FPN_1x.yaml (or other configuration files)

Download here:

BMaskR-CNN/projects/BMaskR-CNN/configs at master · hustvl/BMaskR-CNN (github.com)

3. Problem 3 encountered: KeyError: ‘Non-existent config key: MODEL.PRETRAINED_MODELS’

refer to

KeyError: ‘Non-existent config key: MODEL.PRETRAINED_MODELS’ · Issue #9 · zjhuang22/maskscoring_rcnn (github.com)

In the root directory->tools->train_net.py, add cfg.set_new_allowed(True) to the setup() function

4. Problem 4 encountered: KeyError: “No object named ‘BoundaryROIHeads’ found in ‘ROI_HEADS’ registry!”

You need to copy the bmaskrcnn folder in the BMaskR-CNN folder in the root directory to the root directory

Then add it in the root directory->tools->train_net.py

import bmaskrcnn

5. Problem 5: How to train your own data set

This part can refer to mask r-cnn based on detectron2 training. The training code is similar. Below is a copy of the code I used for training.

#!/usr/bin/env python
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
"""
Detection 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 torch

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
importsys

sys.path.append('BMaskR-CNN')

import bmaskrcnn
from detectron2.data.datasets import register_coco_instances
from detectron2.data import DatasetCatalog, MetadataCatalog
from detectron2.data.datasets.coco import load_coco_json
import pycocotools

# Declare categories, try to keep
CLASS_NAMES = ["_background_", 'Category 1', 'Category 2', 'Category 3, and so on']
#Dataset path
DATASET_ROOT = 'Put your data set root directory here'
ANN_ROOT = os.path.join(DATASET_ROOT, 'annotations')

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

TRAIN_JSON = os.path.join(ANN_ROOT, 'instances_train2014.json')
# VAL_JSON = os.path.join(ANN_ROOT, 'val.json')
VAL_JSON = os.path.join(ANN_ROOT, 'instances_val2014.json')

#Declare a subset of the dataset
PREDEFINED_SPLITS_DATASET = {
    "train2014": (TRAIN_PATH, TRAIN_JSON),
    "val2014": (VAL_PATH, VAL_JSON),
}


# ===========There are two ways to register a data set. The second plain_register_dataset method that I use directly can also be used in the form of register_dataset============ ======
# Register the data set (this step is to register the custom data set into Detectron2)
def register_dataset():
    """
    purpose: register all splits of dataset with PREDEFINED_SPLITS_DATASET
    """
    for key, (image_root, json_file) in PREDEFINED_SPLITS_DATASET.items():
        register_dataset_instances(name=key,
                                   json_file=json_file,
                                   image_root=image_root)


# Register the data set instance and load the object instances in the data set
def register_dataset_instances(name, json_file, image_root):
    """
    purpose: register dataset to DatasetCatalog,
             register metadata to MetadataCatalog and set attribute
    """
    DatasetCatalog.register(name, lambda: load_coco_json(json_file, image_root, name))
    MetadataCatalog.get(name).set(json_file=json_file,
                                  image_root=image_root,
                                  evaluator_type="coco")


# =============================
# Register dataset and metadata
def plain_register_dataset():
    # Training set
    DatasetCatalog.register("train2014", lambda: load_coco_json(TRAIN_JSON, TRAIN_PATH))
    MetadataCatalog.get("train2014").set(thing_classes=CLASS_NAMES, # You can choose to turn it on, but Chinese cannot be displayed. Please note here that it is best to turn it off for Chinese.
                                         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"))
    # Validation/test set
    DatasetCatalog.register("val2014", lambda: load_coco_json(VAL_JSON, VAL_PATH))
    MetadataCatalog.get("val2014").set(thing_classes=CLASS_NAMES, # You can choose to turn it on, but it cannot display Chinese. Please note here that it is best to turn it off for Chinese.
                                       evaluator_type='coco', #Specify the evaluation method
                                       json_file=VAL_JSON,
                                       image_root=VAL_PATH)


# Check the data set annotation and visually check whether the data set annotation is correct.
# You can also write your own script to judge this. In fact, it is to judge whether the label box exceeds the image boundary.
# Optionally use this method
def checkout_dataset_annotation(name="val2014"):
    # 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


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):
        """
        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,
                    num_classes=cfg.MODEL.SEM_SEG_HEAD.NUM_CLASSES,
                    ignore_label=cfg.MODEL.SEM_SEG_HEAD.IGNORE_VALUE,
                    output_dir=output_folder,
                )
            )
        if evaluator_type in ["coco", "coco_panoptic_seg"]:
            evaluator_list.append(COCOEvaluator(dataset_name, cfg, True, output_folder))
        if evaluator_type == "coco_panoptic_seg":
            evaluator_list.append(COCOPanopticEvaluator(dataset_name, output_folder))
        if evaluator_type == "cityscapes_instance":
            assert(
                    torch.cuda.device_count() >= comm.get_rank()
            ), "CityscapesEvaluator currently do not work with multiple machines."
            return CityscapesInstanceEvaluator(dataset_name)
        if evaluator_type == "cityscapes_sem_seg":
            assert(
                    torch.cuda.device_count() >= comm.get_rank()
            ), "CityscapesEvaluator currently do not work with multiple machines."
            return CityscapesSemSegEvaluator(dataset_name)
        elif evaluator_type == "pascal_voc":
            return PascalVOCDetectionEvaluator(dataset_name)
        elif evaluator_type == "lvis":
            return LVISEvaluator(dataset_name, cfg, True, 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)

    @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()
    # args.config_file = "configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x_runway.yaml"
    cfg.set_new_allowed(True)
    cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    # Change configuration parameters
    cfg.DATASETS.TRAIN = ("train2014",) # Training data set name
    cfg.DATASETS.TEST = ("val2014",) # Verification data set name
    cfg.DATALOADER.NUM_WORKERS = 4 # Threads
    cfg.MODEL.RETINANET.NUM_CLASSES = 3 + 1

    cfg.freeze()
    default_setup(cfg, args)
    returncfg


def main(args):
    cfg = setup(args)
    plain_register_dataset() # # Modify

    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()
    # register_coco_instances("bars", {}, "/public/home/chenweiwen/BMaskR-CNN/datasets/coco/annotations/instances_train2014.json", "/public/home/chenweiwen/BMaskR-CNN/ datasets/coco/train2014/")
    # MODEL.ROI_HEADS.NUM_CLASSES = 3
    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,),
    )