深度学习模型库detectron2的安装、使用和自定义数据集训练

一、介绍

detectron2库是由facebookresearch团队开发的开源项目，github地址为https://github.com/facebookresearch/detectron2.git，这里封装集成了很多深度学习的模型，包括基础的rcnn系列以及transform系列，详细了解可以查看官方文档https://detectron2.readthedocs.io/en/latest/index.html，接下来介绍的安装使用训练都是基于官方文档的介绍。
下面的安装使用训练都是基于centos系统的。其它环境请查阅官方文档。

二、安装

安装detectron2库，首先需要检查一下自己的环境，到目前为止，该库只能用在macos和linux系统中，PyTorch ≥ 1.8 并且要安装与PyTorch安装相匹配的 torchvision，OpenCV 是可选的，但演示和可视化需要，建议安装，但是我在使用的过程中opencv会有bug，编译代码会报错，将opencv卸载后重装就可以，环境满足后，将库整个克隆下来，不要进入项目，在整个项目的上级目录里面执行下面的代码：
python -m pip install -e detectron2
等待系统编译完成即可。

三、使用

在该项目中，可以使用已经训练好的网络做演示，只需要进入demo目录：
cd demo/
然后再在命令行中输入：
python demo.py --config-file ../configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml \ --input input1.jpg input2.jpg \ [--other-options] --opts MODEL.WEIGHTS detectron2://COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl
要在您的网络摄像头上运行，请替换为--input files--webcam；
要在视频上运行，请替换为--input files--video-input video.mp4；
要在 cpu 上运行，请在--opts中添加MODEL.DEVICE cpu；
要将输出保存到目录（用于图像）或文件（用于网络摄像头或视频），请使用--output；
不同模型的模型权重文件可以通过库中的model_zoo.get_checkpoint_url()函数获得，在函数中输入需要的模型配置文件作为输入参数即可，例如“COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml”。

四、训练自己的数据集

训练自己的数据集，首先你得有自己的数据集，并且是coco数据集格式，如果只有图片，可以用labelme程序自行标注，获得json格式的标注文件，再通过github上的labelme2coco项目整合含有标注的json文件，获得coco格式的数据集。
训练项目的目录大致如下：

其中test目录下有测试数据集图片，train目录下为寻来你数据集图片，test.json、train.json
分别为测试图片集和训练图片集的标注文件，完成utils.py和train.py的编程即可开始训练自己的数据集（要在含有detectron2库的环境下训练），
其中utils.py文件中的代码如下：

from detectron2.data import DatasetCatalog,MetadataCatalog
from detectron2.utils.visualizer import Visualizer
from detectron2.config import get_cfg
from detectron2 import model_zoo

from detectron2.utils.visualizer import ColorMode

import random
import cv2
import matplotlib.pyplot as plt

def plot_samples(dataset_name,n=1):
    dataset_custom=DatasetCatalog.get(dataset_name)
    dataset_custom_metadata=MetadataCatalog.get(dataset_name)
    for s in random.sample(dataset_custom,n):
        img=cv2.imread(s["file_name"])
        v=Visualizer(img[:,:,::-1],metadata=dataset_custom_metadata,scale=0.5)
        v=v.draw_dataset_dict(s)
        plt.figure(figsize=(15,20))
        plt.imshow(v.get_image())
        plt.savefig('./outputs/sample.png')
        plt.show()

def get_train_cfg(config_file_path,checkpoint_url,train_dataset_name,test_dataset_name,num_classes,device,output_dir):
    cfg=get_cfg()

    cfg.merge_from_file(model_zoo.get_config_file(config_file_path))
    cfg.MODEL.WEIGHTS=model_zoo.get_checkpoint_url(checkpoint_url)
    cfg.DATASETS.TRAIN=(train_dataset_name,)
    cfg.DATASETS.TEST=(test_dataset_name,)

    cfg.DATALOADER.NUM_WORKERS=2

    cfg.SOLVER.IMS_PER_BATCH=2
    cfg.SOLVER.BASE_LR=0.00025
    cfg.SOLVER.MAX_ITER=1000
    cfg.SOLVER.STEPS=[]

    cfg.MODEL.ROI_HEADS.NUM_CLASSES=num_classes
    cfg.MODEL.DEVICE=device
    cfg.OUTPUT_DIR=output_dir

    return cfg

def on_image(image_path,predictor):
    im=cv2.imread(image_path)
    outputs=predictor(im)
    v=Visualizer(im[:,:,::-1],metadata={},scale=0.5,instance_mode=ColorMode.SEGMENTATION)
    v=v.draw_instance_predictions(outputs["instance"].to("cpu"))

    plt.Figure(figsize=(14,10))
    plt.imshow(v.get_image())
    plt.savefig('predict')

# def on_video(videoPath,predictor):
#     cap=cv2.VideoCapture(videoPath)
#     if(cap.isOpened()==False):
#         peint("Error opening file...")
#         return
#
#     (success,image)=cap.read()
#     while success:
#         predictions=predictor(image)
#         v=Visualizer(image[:,:,::-1],metadata={},instance_model=ColorMode.SEGMENTATION)
#         output=v.draw_instance_predictions(predictions["instances"].to("cpu"))
#
#         cv2.imshow("Result",output.get_image()[:,:,::-1])
#
#         key=cv2.waitKey(1)&0xFF
#         if key==ord("q"):
#             break
#         (success,image)=cap.read()

train.py文件中的代码如下：

from detectron2.utils.logger import setup_logger
setup_logger()

from detectron2.data.datasets import register_coco_instances
from detectron2.engine import DefaultTrainer

import os
import pickle
from utils import *

config_file_path="COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"
checkpoint_url="COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"

ouput_dir="./output"
num_classes=1
device="cpu" #cuda

train_dataset_name="LP_train"
train_image_path="train"
train_json_annot_path="train.json"

test_dataset_name="LP_test"
test_image_path="test"
test_json_annot_path="test.json"

cfg_save_path="OD_cfg.pickle"

register_coco_instances(name=train_dataset_name,metadata={},json_file=train_json_annot_path,image_root=train_image_path)
register_coco_instances(name=test_dataset_name,metadata={},json_file=test_json_annot_path,image_root=test_image_path)

# plot_samples(dataset_name=train_dataset_name,n=2)

def main():
    cfg=get_train_cfg(config_file_path,checkpoint_url,train_dataset_name,test_dataset_name,num_classes,device,ouput_dir)
    with open(cfg_save_path,'wb') as f:
        pickle.dump(cfg,f,protocol=pickle.HIGHEST_PROTOCOL)

    os.makedirs(cfg.OUTPUT_DIR,exist_ok=True)

    trainer=DefaultTrainer(cfg)
    trainer.resume_or_load(resume=False)

    trainer.train()

if __name__=='__main__':
    main()

test.py文件中的代码如下：

from detectron2.engine import DefaultPredictor

import os
import pickle
from utils import *

cfg_save_path="OD_cfg.pickle"

with open(cfg_save_path,'rb') as f:
    cfg=pickle.load(f)

cfg.MODEL.WEIGHTS=os.path.join(cfg.OUTPUT_DIR,"model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST=0.5

predictor=DefaultPredictor(cfg)
image_path="test/20170510_141006.jpg"
#video_path="test/abc.mp4"

on_image(image_path,predictor)