detectron2安装（ubuntu）及训练自定义数据集（maskrcnn）

目录

环境搭建

测试demo

自定义数据集注册

训练自己的数据集

推理图片

常见问题

---

环境搭建

1、创建虚拟环境

conda create -n det-maskrcnn python=3.7

2、进入环境安装依赖包

conda activate det-maskrcnn    # 进入虚拟环境
pip install opencv-python matplotlib cython tensorboardX   
conda install pytorch==1.7.0 torchvision==0.8.0 torchaudio==0.7.0 cudatoolkit=10.1 -c pytorch  # 根据自己电脑情况到官网https://pytorch.org/安装
# 需要注意detectron2不同版本对pytorch版本要求不同，安装前可先查看detectron2的setup.py文件

下载cocoapi放到项目目录下
cd cocoapi/PythonAPI
python setup.py build_ext install

安装detectron2，回到detectron2根目录
python -m pip install -e .

可能出现问题：

报错：ERROR: Command errored out with exit status 1: /root/miniconda3/envs/det-maskrcnn/bin/python -c 'import io, os, sys, setuptools, tokenize; sys.argv[0] = '"'"'/root/workspace/maskrcnn/detectron2-maskrcnn/setup.py'"'"'; __file__='"'"'/root/workspace/maskrcnn/detectron2-maskrcnn/setup.py'"'"';f = getattr(tokenize, '"'"'open'"'"', open)(__file__) if os.path.exists(__file__) else io.StringIO('"'"'from setuptools import setup; setup()'"'"');code = f.read().replace('"'"'\r\n'"'"', '"'"'\n'"'"');f.close();exec(compile(code, __file__, '"'"'exec'"'"'))' develop --no-deps Check the logs for full command output.

原因：gcc版本有问题，卸载重新安装

---

测试demo

python demo/demo.py --config-file configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml --input demo/demo.jpg --opts MODEL.WEIGHTS weights/model_final_r50_fpn_1x.pkl

预训练模型可以通过官方的 MODEL_ZOO 下载。

---

自定义数据集注册

1、将自己的数据集转换为coco格式，并按照以下格式放置数据集

datasets
   -| coco
      -| annotations      # 存放注释文件
          -| instances_train2017.json
          -| instances_val2017.json
      -| train2017        # 存放训练图片
      -| val2017          # 存放验证图片

2、在train_net.py的基础上，对其进行更改并另存为train.py文件

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.datasets import load_coco_json
from detectron2.data import MetadataCatalog, DatasetCatalog
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.datasets import register_coco_instances


# 单GPU
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
#声明类别，尽量保持
CLASS_NAMES = ["bv", "rb"]
DATASET_CATEGORIES = [
    {"name": "bv", "id": 1, "isthing": 1, "color": [220, 20, 60]},
    {"name": "rb", "id": 2, "isthing": 1, "color": [219, 142, 185]},
]
# 声明数据集的子集
PREDEFINED_SPLITS_DATASET = {
    "coco_my_train": ("./datasets/coco/train2017", "./datasets/coco/annotations/instances_train2017.json"),
    "coco_my_val": ("./datasets/coco/val2017", "./datasets/coco/annotations/instances_val2017.json"),
}
# 注册数据集
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,
                                   metadate=get_dataset_instances_meta(),