detectron2 配置与使用

下载与配置

github地址 https://github.com/facebookresearch/detectron2
官方文档 https://detectron2.readthedocs.io/en/latest/index.html

环境
python >=3.6
pytorch >=1.7
torchvision
opencv （训练不需要，但使用demo和显示图片需要）

python setup.py build develop

数据集注册和训练的准备

在detectron2/tools/train_net.py文件中插入代码：

from detectron2.data.datasets import register_coco_instances

register_coco_instances("coco_my_train", {},
                         "/data/输入训练图片的json文件/train.json",
                         "/data/输入训练图片文件夹/train")
register_coco_instances("coco_my_val", {},
                         "/data/输入验证图片的json文件/val.json",
                         "/data/输入验证图片文件夹/val")
CLASS_NAMES =["__background__","class1",,,,,,"class n"]
#若json文件中的类别是从0开始，则不需要开头的background，若是从1开始，则需要

在此文件中的setup函数中，可以配置超参数：

	#选择模型
    args.config_file = "../configs/COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml" 
    # 预训练模型权重,pkl,pth文件均可，也可以在模型文件中直接修改
    cfg.MODEL.WEIGHTS = "/home/detectron2/modelweight.pkl"  
	cfg.DATASETS.TRAIN = ("coco_my_train",)  # 训练数据集名称
    cfg.DATASETS.TEST = ("coco_my_val",)
    cfg.DATALOADER.NUM_WORKERS = 4  # 单线程
    cfg.INPUT.CROP.ENABLED = True
    # x为文件中目标的类别数，类别数+1因为有background，若没有则不加
    cfg.MODEL.RETINANET.NUM_CLASSES = x+1
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = x+1
    

    cfg.SOLVER.IMS_PER_BATCH = 2  # batch_size=2
    # 根据训练数据总数目以及batch_size，计算出每个epoch需要的迭代次数
    # 7000为你的训练集中图片数据的总数
    ITERS_IN_ONE_EPOCH =  7000
    # 指定最大迭代次数       
    cfg.SOLVER.MAX_ITER =   (ITERS_IN_ONE_EPOCH * 5) - 1  # 5个epoch
    # 初始学习率
    cfg.SOLVER.BASE_LR = 0.002
    # 优化器动能
    cfg.SOLVER.MOMENTUM = 0.9
    # 权重衰减
    cfg.SOLVER.WEIGHT_DECAY = 0.0001
    cfg.SOLVER.WEIGHT_DECAY_NORM = 0.0
    # 学习率衰减倍数
    cfg.SOLVER.GAMMA = 0.1
    # 迭代到指定次数，学习率进行衰减
    cfg.SOLVER.STEPS = (7000,)
    # 在训练之前，会做一个热身运动，学习率慢慢增加初始学习率
    cfg.SOLVER.WARMUP_FACTOR = 1.0 / 1000
    # 热身迭代次数
    cfg.SOLVER.WARMUP_ITERS = 1000
    cfg.SOLVER.WARMUP_METHOD = "linear"
    # 保存模型文件的命名数据减1
    cfg.SOLVER.CHECKPOINT_PERIOD = ITERS_IN_ONE_EPOCH - 1
    # 迭代到指定次数，进行一次评估
    cfg.TEST.EVAL_PERIOD = ITERS_IN_ONE_EPOCH

    cfg.INPUT.CROP.ENABLED = True
    cfg.INPUT.MAX_SIZE_TRAIN = 1440 # 训练图片输入的最大尺寸
    cfg.INPUT.MAX_SIZE_TEST = 1440 # 测试数据输入的最大尺寸
    cfg.INPUT.MIN_SIZE_TRAIN = 1024 # 训练图片输入的最小尺寸，可以设定为多尺度训练
    cfg.INPUT.MIN_SIZE_TEST = 1024
    #cfg.INPUT.MIN_SIZE_TRAIN_SAMPLING，其存在两种配置，分别为 choice 与 range ：
    # range 让图像的短边从 512-768随机选择
    #choice ： 把输入图像转化为指定的，有限的几种图片大小进行训练，即短边只能为 512或者768
    cfg.INPUT.MIN_SIZE_TRAIN_SAMPLING = 'range'

修改 detectron2/detectron2/data/datasets/builtin_meta.py文件
将COCO_CATEGORIES 中的类别改为自己数据集的类别，注意background。

COCO_CATEGORIES = [
    {"color": [220, 20, 60], "isthing": 0, "id": 0, "name": "__background__"},
    {"color": [119, 11, 32], "isthing": 1, "id": 1, "name": "dog"},
    {"color": [0, 0, 142], "isthing": 1, "id": 2, "name": "cat"},

在 _get_coco_instances_meta函数中修改

assert len(thing_ids) == x, len(thing_ids)
#                        ↑ 把x改为自己数据集的类别数

在_get_coco_panoptic_separated_meta函数中

    assert len(stuff_ids) == y, len(stuff_ids)
    #						#↑ 背景类数目（有没有background），有1无0

训练文件的修改和配置暂时就这么多。

Training

python train.py
参数基本上在setup中都配置好了，一般只需要 – num-gpus 选择使用gpu数目和 --resume 继续上次的训练

查看标注好的图片

detectron2目录下创建一个新.py文件，使用以下代码

import os
import time
import cv2
import numpy as np
from detectron2 import model_zoo
from detectron2.config import get_cfg
from detectron2.data import DatasetCatalog
from detectron2.data import MetadataCatalog
from detectron2.data.datasets import register_coco_instances
from detectron2.engine import DefaultPredictor
from detectron2.utils.visualizer import ColorMode, Visualizer

if __name__ == '__main__':
    cfg = get_cfg()
    cfg.MODEL.DEVICE = 'cpu'     # 这里可定义是调用CPU还是GPU，默认GPU
    cfg.OUTPUT_DIR = '../output'	
    register_coco_instances("coco_my_test", {},
		"/data/train.json",
		"/data/test")
    metadata = MetadataCatalog.get("coco_my_test")
    cfg.DATALOADER.NUM_WORKERS = 4  # 单线程
    cfg.INPUT.CROP.ENABLED = True
    #训练好的权重文件放在模型文件中了。
    cfg.merge_from_file("./configs/COCO-Detection/faster_rcnn_R_50_C4_1x.yaml")
    cfg.DATASETS.TRAIN = ("coco_my_test",)
    cfg.DATASETS.TEST = ()

    #cfg.MODEL.WEIGHTS = 
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = x # 类别数量，此参数若是不能与训练时的保持一致，测试结果将会严重错误（如：类别错位，只显示一种类别等）
    score_test = 0.5
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = score_test  # 低于50分的box过滤掉
    dataset_dicts = DatasetCatalog.get("coco_my_test")
    # 测试图片保存路径
    path_saveFile = "../need_test_img/"
    # 判断保存目录路径是否存在，若不存在则创建
    if os.path.isdir(path_saveFile):
        pass
    else:
        os.mkdir(path_saveFile)
    # 判断要测试的子文件是否存在，若不存在则创建
    path_saveImg = path_saveFile + "/test_project/"
    if os.path.isdir(path_saveImg):
        pass
    else:
        os.mkdir(path_saveImg)

    # 遍历所有待检测图片
    print("\n\nPlease wait a moment...")
    for root, dirs, files in os.walk("../need_test_img/"):
        if len(files) == 0:
            print("The file has no image...")
            break
        for d in dirs:
            print("image number =", d)
        for file in files:
            image_path = root + '/' + file
            img_name = image_path.split('/')[-1].split('.')[0]
            img = cv2.imread(image_path)
            imgH, imgW = img.shape[:2]
            maxValue = max(imgH, imgW)

            # cfg.INPUT.MAX_SIZE_TEST = int(maxValue/2.66)
            # cfg.INPUT.MIN_SIZE_TEST = int(maxValue/1.25)
            predictor = DefaultPredictor(cfg)

            outputs_test = predictor(img) 
            #print("outputs_test =", outputs_test)
            v_test = Visualizer(img[:, :, ::-1],
                                metadata=metadata,
                                scale=1,
                                instance_mode=ColorMode.IMAGE_BW)
            # print("v_test =", v_test)
            out_test = v_test.draw_instance_predictions(
                outputs_test["instances"].to("cpu"))

            # 写入检测好的图片
            cv2.imwrite(path_saveImg + img_name + ".jpg",
                        out_test.get_image()[:, :, ::-1])
    print("\n\nEnd of test process...")

python show.py
就直接在输出文件夹中得到检测结果了。