跑bmask-rcnn遇到的几个小坑

源码地址

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

1.遇到的问题1：安装detectron2

请务必按照官方教程安装，尤其是使用windows+anaconda，推荐使用源码安装

切勿pip

官方安装教程：

Use Custom Datasets — detectron2 0.6 documentation

2.遇到的问题2：缺少文件bmask_rcnn_R_50_FPN_1x.yaml（或者其他配置文件）

在这里下载：

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

3.遇到的问题3：KeyError: 'Non-existent config key: MODEL.PRETRAINED_MODELS'

参考

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

在根目录->tools->train_net.py中，setup()函数中添加cfg.set_new_allowed(True)

4.遇到的问题4：KeyError: "No object named 'BoundaryROIHeads' found in 'ROI_HEADS' registry!"

需要把根目录下的BMaskR-CNN文件夹内的bmaskrcnn文件夹，复制到根目录下

然后在根目录->tools->train_net.py中添加

import bmaskrcnn

5.遇到的问题5：如何训练自己的数据集

此部分可参考基于detectron2训练的mask r-cnn，训练代码是类似的，下面贴了一份我用来训练的代码

#!/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
import sys

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

# 声明类别，尽量保持
CLASS_NAMES = ["_background_", '类别1', '类别2', '类别3，以此类推']
# 数据集路径
DATASET_ROOT = '这里放你的数据集根目录'
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')

# 声明数据集的子集
PREDEFINED_SPLITS_DATASET = {
    "train2014": (TRAIN_PATH, TRAIN_JSON),
    "val2014": (VAL_PATH, VAL_JSON),
}


# ===========以下有两种注册数据集的方法，本人直接用的第二个plain_register_dataset的方式 也可以用register_dataset的形式==================
# 注册数据集（这一步就是将自定义数据集注册进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)


# 注册数据集实例，加载数据集中的对象实例
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")


# =============================
# 注册数据集和元数据
def plain_register_dataset():
    # 训练集
    DatasetCatalog.register("train2014", lambda: load_coco_json(TRAIN_JSON, TRAIN_PATH))
    MetadataCatalog.get("train2014").set(thing_classes=CLASS_NAMES,  # 可以选择开启，但是不能显示中文，这里需要注意，中文的话最好关闭
                                         evaluator_type='coco',  # 指定评估方式
                                         json_file=TRAIN_JSON,
                                         image_root=TRAIN_PATH)

    # DatasetCatalog.register("coco_my_val", lambda: load_coco_json(VAL_JSON, VAL_PATH, "coco_2017_val"))
    # 验证/测试集
    DatasetCatalog.register("val2014", lambda: load_coco_json(VAL_JSON, VAL_PATH))
    MetadataCatalog.get("val2014").set(thing_classes=CLASS_NAMES,  # 可以选择开启，但是不能显示中文，这里需要注意，中文的话最好关闭
                                       evaluator_type='coco',  # 指定评估方式
                                       json_file=VAL_JSON,
                                       image_root=VAL_PATH)


# 查看数据集标注，可视化检查数据集标注是否正确，
# 这个也可以自己写脚本判断，其实就是判断标注框是否超越图像边界
# 可选择使用此方法
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)
    # 更改配置参数
    cfg.DATASETS.TRAIN = ("train2014",)  # 训练数据集名称
    cfg.DATASETS.TEST = ("val2014",)  # 验证数据集名称
    cfg.DATALOADER.NUM_WORKERS = 4  # 线程
    cfg.MODEL.RETINANET.NUM_CLASSES = 3 + 1

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


def main(args):
    cfg = setup(args)
    plain_register_dataset()  # # 修改

    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,),
    )