本教程实现使用Detectron2库训练自定义目标检测器。值得注意的是，Detectron2库远不止目标检测，它支持语义分割、关键点检测、掩码和稠密估计。
我们将介绍在自定义对象上训练Detectron2所需的步骤。我们使用公共血细胞检测数据集。
为了训练我们的检测器，我们采取以下步骤:

• 获取Detectron2的Docker镜像，并创建自己的容器
• 可视化Detectron2的训练数据
• 编写我们的Detectron2训练配置
• 运行Detectron2训练
• 评估Detectron2性能
• 在测试图像上运行Detectron2推理

1.获取Detectron2的Docker镜像，并创建自己的容器

(1) 拉取Docker容器

docker pull mendelxu/mmdet:nvcr_21_05_detectron2

(2)创建容器

docker run -it -d --gpus all --restart always --name xw_dev_detectron2 -p 10005:22 -v  /home/xx/xw:/xw --ipc=host mendelxu/mmdet:nvcr_21_05_detectron2 /bin/bash

2.导入相关库

import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()

# import some common libraries
import numpy as np
import cv2
import random

# import some common detectron2 utilities
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
from detectron2.data.catalog import DatasetCatalog
import os
# 使用4，5两块GPU
os.environ["CUDA_VISIBLE_DEVICES"] = "4,5"

3.导入和注册自定义Detectron2数据

from detectron2.data.datasets import register_coco_instances
register_coco_instances("my_dataset_train", {}, "/xw/datasets/BCCD/416x416_aug/BCCD.v4-416x416_aug.coco/train/_annotations.coco.json", "/xw/datasets/BCCD/416x416_aug/BCCD.v4-416x416_aug.coco/train")
register_coco_instances("my_dataset_val", {}, "/xw/datasets/BCCD/416x416_aug/BCCD.v4-416x416_aug.coco/valid/_annotations.coco.json", "/xw/datasets/BCCD/416x416_aug/BCCD.v4-416x416_aug.coco/valid")
register_coco_instances("my_dataset_test", {}, "/xw/datasets/BCCD/416x416_aug/BCCD.v4-416x416_aug.coco/test/_annotations.coco.json", "/xw/datasets/BCCD/416x416_aug/BCCD.v4-416x416_aug.coco/test")

4.可视化训练数据集

my_dataset_train_metadata = MetadataCatalog.get("my_dataset_train")
dataset_dicts = DatasetCatalog.get("my_dataset_train")

import random
from detectron2.utils.visualizer import Visualizer

for idx, d in enumerate(random.sample(dataset_dicts, 3)):
    img = cv2.imread(d["file_name"])
    visualizer = Visualizer(img[:, :, ::-1], metadata=my_dataset_train_metadata, scale=0.5)
    vis = visualizer.draw_dataset_dict(d)
    img = vis.get_image()[:, :, ::-1]
    cv2.imwrite(str(idx) + ".jpg", img)

5.训练自定义的Detectron2检测器

from detectron2.engine import DefaultTrainer
from detectron2.evaluation import COCOEvaluator
import os
# 在开始训练之前，我们需要确保模型根据我们的验证集进行验证。不幸的是，这在默认情况下不会发生🤔。我们需要如下类：
class CocoTrainer(DefaultTrainer):

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

    if output_folder is None:
        os.makedirs("coco_eval", exist_ok=True)
        output_folder = "coco_eval"

    return COCOEvaluator(dataset_name, cfg, False, output_folder)

from detectron2.config import get_cfg
#from detectron2.evaluation.coco_evaluation import COCOEvaluator
import os

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("my_dataset_train",)
cfg.DATASETS.TEST = ("my_dataset_val",)

cfg.DATALOADER.NUM_WORKERS = 4
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml")  # Let training initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 4
cfg.SOLVER.BASE_LR = 0.001


cfg.SOLVER.WARMUP_ITERS = 1000
cfg.SOLVER.MAX_ITER = 1500 #adjust up if val mAP is still rising, adjust down if overfit
cfg.SOLVER.STEPS = (1000, 1200)
cfg.SOLVER.GAMMA = 0.05

cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 64
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 4 #your number of classes + 1

cfg.TEST.EVAL_PERIOD = 500

os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = CocoTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()

我们在这里调用的模型是Faster RCNN。Detectron2允许您使用许多模型架构选项，您可以在Detectron2模型库中看到这些选项。
我们所做的另一个配置选择是MAX_ITER参数。这指定了模型将训练多长时间，您可能数据集进行调整。

6.测试集评估

#test evaluation
from detectron2.data import DatasetCatalog, MetadataCatalog, build_detection_test_loader
from detectron2.evaluation import COCOEvaluator, inference_on_dataset

cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.85
predictor = DefaultPredictor(cfg)
evaluator = COCOEvaluator("my_dataset_test", cfg, False, output_dir="./output/")
val_loader = build_detection_test_loader(cfg, "my_dataset_test")
inference_on_dataset(trainer.model, val_loader, evaluator)

7.在测试图像上运行Detectron2推理

cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.DATASETS.TEST = ("my_dataset_test", )
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7   # set the testing threshold for this model
predictor = DefaultPredictor(cfg)
test_metadata = MetadataCatalog.get("my_dataset_test")
from detectron2.utils.visualizer import ColorMode
import glob

for idx, imageName in enumerate(glob.glob('/xw/datasets/BCCD/416x416_aug/BCCD.v4-416x416_aug.coco/test/*jpg')):
    if idx > 2:
        break
    im = cv2.imread(imageName)
    outputs = predictor(im)
    v = Visualizer(im[:, :, ::-1],
                metadata=test_metadata, 
                scale=0.8
                    )
    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    cv2.imwrite("infer_"+str(idx)+".jpg", out.get_image()[:, :, ::-1])