本教程实现使用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])