Detectron2 官方文檔里的 Getting Started 提供了兩種使用 detectron2 的樣例。其一是讀者大概率已經閱讀過的 Colab Notebook ——騎馬王子和氣球檢測，其二是使用命令行執行的 python 文件，包括演示文件 demo.py 及自行用於部署的 train_net.py & plain_train_net.py 。Notebook 已述明使用 Mask-RCNN 進行 mask detection 的簡單步驟，包括注冊數據集、配置 config 以及訓練和驗證、可視化結果。由於人力標注費力的原因，自行標注的數據集通常只有目標物體 box 而未標注 mask，因此需使用 Faster-RCNN 進行訓練，有關代碼可參考 Medium 文章 以及其 代碼。
本篇參考 文章 ，簡述 python 文件 + 命令行執行的項目部署方法，以便於在 GPU 服務器端運行。
目錄一覽
本項目部署目錄如下：
Project
--configs
----COCO-Detection
faster_rcnn_R_50_FPN_3x.yaml # 官方代碼庫 copy
Base-RCNN-FPN.yaml # 同上
my_config.yaml # 自己的訓練配置文件
--tools
train_net.py # 官方代碼庫 copy 並自行修改
train.sh
train_resume.sh
eval.sh
--utils
txt2coco.py # 將自己的數據集轉化為標准 coco 數據集
數據集目錄如下：
MyDataset
--train
01_00001.jpg # 圖片名無所謂
...
--val
...
train.json # 位置自行決定
val.json
數據集注冊
使用 CocoFormat 的數據集是最優雅的做法，當然你也可以按照官方給定的方法在 tools/rain_net.py 中自定義數據集。關於 csv/voc/labelme 等格式向 coco 格式的轉化，可以參考 Github 的代碼。本文 utils/txt2coco.py 基於上述代碼將自行標注的 txt 格式數據轉化為了 coco 格式。
為方便數據集注冊，本文借鑒 文章 ，將數據集注冊代碼包裝成類：
from detectron2.data import DatasetCatalog, MetadataCatalog
from detectron2.data.datasets.coco import load_coco_json
class Register:
"""用於注冊自己的數據集"""
CLASS_NAMES = ['__background__', 'Red', 'Blue', 'Yellow', 'White', 'Black', 'Other', 'NoHelmet'] # 保留 background 類
ROOT = "/home/your_dataset_dir"
def __init__(self):
self.CLASS_NAMES = Register.CLASS_NAMES or ['__background__', ]
# 數據集路徑
self.DATASET_ROOT = Register.ROOT or '/home/yourdir'
# ANN_ROOT = os.path.join(self.DATASET_ROOT, 'COCOformat')
self.ANN_ROOT = self.DATASET_ROOT
self.TRAIN_PATH = os.path.join(self.DATASET_ROOT, 'train')
self.VAL_PATH = os.path.join(self.DATASET_ROOT, 'val')
self.TRAIN_JSON = os.path.join(self.ANN_ROOT, 'train.json')
self.VAL_JSON = os.path.join(self.ANN_ROOT, 'val.json')
# VAL_JSON = os.path.join(self.ANN_ROOT, 'test.json')
# 聲明數據集的子集
self.PREDEFINED_SPLITS_DATASET = {
"coco_my_train": (self.TRAIN_PATH, self.TRAIN_JSON),
"coco_my_val": (self.VAL_PATH, self.VAL_JSON),
}
def register_dataset(self):
"""
purpose: register all splits of datasets with PREDEFINED_SPLITS_DATASET
注冊數據集（這一步就是將自定義數據集注冊進Detectron2）
"""
for key, (image_root, json_file) in self.PREDEFINED_SPLITS_DATASET.items():
self.register_dataset_instances(name=key,
json_file=json_file,
image_root=image_root)
@staticmethod
def register_dataset_instances(name, json_file, image_root):
"""
purpose: register datasets 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(self):
"""注冊數據集和元數據"""
# 訓練集
DatasetCatalog.register("coco_my_train", lambda: load_coco_json(self.TRAIN_JSON, self.TRAIN_PATH))
MetadataCatalog.get("coco_my_train").set(thing_classes=self.CLASS_NAMES, # 可以選擇開啟，但是不能顯示中文，這里需要注意，中文的話最好關閉
evaluator_type='coco', # 指定評估方式
json_file=self.TRAIN_JSON,
image_root=self.TRAIN_PATH)
# DatasetCatalog.register("coco_my_val", lambda: load_coco_json(VAL_JSON, VAL_PATH, "coco_2017_val"))
# 驗證/測試集
DatasetCatalog.register("coco_my_val", lambda: load_coco_json(self.VAL_JSON, self.VAL_PATH))
MetadataCatalog.get("coco_my_val").set(thing_classes=self.CLASS_NAMES, # 可以選擇開啟，但是不能顯示中文，這里需要注意，中文的話最好關閉
evaluator_type='coco', # 指定評估方式
json_file=self.VAL_JSON,
image_root=self.VAL_PATH)
def checkout_dataset_annotation(self, name="coco_my_val"):
"""
查看數據集標注，可視化檢查數據集標注是否正確，
這個也可以自己寫腳本判斷，其實就是判斷標注框是否超越圖像邊界
可選擇使用此方法
"""
# dataset_dicts = load_coco_json(TRAIN_JSON, TRAIN_PATH, name)
dataset_dicts = load_coco_json(self.TRAIN_JSON, self.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
將以上類置於 tools/train_net.py ，並在 main() 函數中第2行添加調用：
def main(args):
cfg = setup(args)
Register().register_dataset() # register my dataset
...
編輯配置文件
從官方下載的 Faster_RCNN 配置文件及其依賴配置最好不要修改，而是基於其配置另寫一個文件 configs/my_config.yaml ：
_BASE_: "COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"
DATASETS:
TRAIN: ("coco_my_train",)
TEST: ("coco_my_val",)
MODEL:
RETINANET:
NUM_CLASSES: 8 # 類別數+1, 因為有background
# WEIGHTS: "../tools/output/model_final.pth"
SOLVER:
# IMS_PER_BATCH: 16
# 初始學習率
BASE_LR: 0.00025
# 迭代到指定次數，學習率進行衰減
# STEPS: (210000, 250000)
# MAX_ITER: 270000
CHECKPOINT_PERIOD: 1000
TEST:
EVAL_PERIOD: 3000
如需修改配置，在該文件中修改就好了。
Train & Eval
控制訓練和測試的命令行代碼置於 shell 文件中更優雅和容易控制。
################ train.sh ################
# Linux 下換行符為 CRLF 的需改為 LF
# lr = 0.00025 * num_gpus
python3 train_net.py \
--config-file ../configs/my_config.yaml \
--num-gpus 4 \
SOLVER.IMS_PER_BATCH 16 \
SOLVER.BASE_LR 0.001 \
SOLVER.MAX_ITER 30000 \
SOLVER.STEPS '(24000, 29000)'
############# train_resume.sh #############
# 斷點續 train
# --num-gpus 親測不能省略
python3 train_net.py \
--config-file ../configs/my_config.yaml \
--num-gpus 4 \
--resume
################# eval.sh #################
python3 train_net.py \
--config-file ../configs/my_config.yaml \
--eval-only \
MODEL.WEIGHTS output/model_final.pth
在命令行中執行 sh train.sh 即可開始訓練。
附錄
自己的 train_net.py 代碼：
#!/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 cv2
import torch
import detectron2.utils.comm as comm
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.config import get_cfg
from detectron2.data import MetadataCatalog, DatasetCatalog
from detectron2.data.datasets import load_coco_json
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.utils.visualizer import Visualizer
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 datasets.
This uses the special metadata "evaluator_type" associated with each builtin datasets.
For your own datasets, 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 datasets {} 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
class Register:
"""用於注冊自己的數據集"""
CLASS_NAMES = ['__background__', 'Red', 'Blue', 'Yellow', 'White', 'Black', 'Other', 'NoHelmet']
ROOT = "/home/mydir"
def __init__(self):
self.CLASS_NAMES = Register.CLASS_NAMES or ['__background__', ]
# 數據集路徑
self.DATASET_ROOT = Register.ROOT or '/home/yourdir'
# ANN_ROOT = os.path.join(self.DATASET_ROOT, 'COCOformat')
self.ANN_ROOT = self.DATASET_ROOT
self.TRAIN_PATH = os.path.join(self.DATASET_ROOT, 'train')
self.VAL_PATH = os.path.join(self.DATASET_ROOT, 'val')
self.TRAIN_JSON = os.path.join(self.ANN_ROOT, 'train.json')
self.VAL_JSON = os.path.join(self.ANN_ROOT, 'val.json')
# VAL_JSON = os.path.join(self.ANN_ROOT, 'test.json')
# 聲明數據集的子集
self.PREDEFINED_SPLITS_DATASET = {
"coco_my_train": (self.TRAIN_PATH, self.TRAIN_JSON),
"coco_my_val": (self.VAL_PATH, self.VAL_JSON),
}
def register_dataset(self):
"""
purpose: register all splits of datasets with PREDEFINED_SPLITS_DATASET
注冊數據集（這一步就是將自定義數據集注冊進Detectron2）
"""
for key, (image_root, json_file) in self.PREDEFINED_SPLITS_DATASET.items():
self.register_dataset_instances(name=key,
json_file=json_file,
image_root=image_root)
@staticmethod
def register_dataset_instances(name, json_file, image_root):
"""
purpose: register datasets 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(self):
"""注冊數據集和元數據"""
# 訓練集
DatasetCatalog.register("coco_my_train", lambda: load_coco_json(self.TRAIN_JSON, self.TRAIN_PATH))
MetadataCatalog.get("coco_my_train").set(thing_classes=self.CLASS_NAMES, # 可以選擇開啟，但是不能顯示中文，這里需要注意，中文的話最好關閉
evaluator_type='coco', # 指定評估方式
json_file=self.TRAIN_JSON,
image_root=self.TRAIN_PATH)
# DatasetCatalog.register("coco_my_val", lambda: load_coco_json(VAL_JSON, VAL_PATH, "coco_2017_val"))
# 驗證/測試集
DatasetCatalog.register("coco_my_val", lambda: load_coco_json(self.VAL_JSON, self.VAL_PATH))
MetadataCatalog.get("coco_my_val").set(thing_classes=self.CLASS_NAMES, # 可以選擇開啟，但是不能顯示中文，這里需要注意，中文的話最好關閉
evaluator_type='coco', # 指定評估方式
json_file=self.VAL_JSON,
image_root=self.VAL_PATH)
def checkout_dataset_annotation(self, name="coco_my_val"):
"""
查看數據集標注，可視化檢查數據集標注是否正確，
這個也可以自己寫腳本判斷，其實就是判斷標注框是否超越圖像邊界
可選擇使用此方法
"""
# dataset_dicts = load_coco_json(TRAIN_JSON, TRAIN_PATH, name)
dataset_dicts = load_coco_json(self.TRAIN_JSON, self.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
def setup(args):
"""
Create configs and perform basic setups.
"""
cfg = get_cfg()
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.freeze()
default_setup(cfg, args)
return cfg
def main(args):
cfg = setup(args)
Register().register_dataset() # register my 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()
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,),
)
參考
[1]Getting Started with Detectron2 - Github
[2]利用detectron2快速使用faster RCNN - Medium
[3]Detectron2訓練自己的數據集（較詳細） - CSDN