深度学习：利用Detectron2训练自己的数据集之目标检测

本文主要阐述如何利用Detectron2来进行目标检测。

一、Detectron2简介

• Detectron2 前身就是鼎鼎大名的Detectron，其实Detectron可以说是Facebook第一代检测工具箱，目前在github上获得 22.5k star！其不仅支持 Detectron已有的目标检测、实例分割、姿态估计等任务，还支持语义分割和全景分割。新增了Cascade R-CNN，Panoptic FPN和TensorMask新模型。
• Detectron2 与mmDetection、SimpleDet被称为目标检测三大神器。

二、基本开发环境搭建

1. 本人使用的开发环境如下：

• 系统：ubuntu20.04(64位)，Python3.7.9；
• 框架：Pytorch1.7.1；
• 图像处理：opencv3.4.2，opencv-python4.5.1.48；
  ubuntu系统环境安装（包括：显卡驱动安装、cuda和cudnn安装）教程有很多，本文不再给出。

2. 安装依赖：

# 创建环境
conda create -n python37_rw python=3.7.9
# 安装 pytorch cpu
conda install pytorch torchvision torchaudio cpuonly -c pytorch
# 安装 pytorch gpu
conda install pytorch torchvision torchaudio cudatoolkit=11.2 -c pytorch
# 安装 opencv
conda install opencv-python
# 安装 flask
conda install flask
# 安装 gcc g++
sudo apt install gcc g++
# 安装 detectron2
python -m pip install -e detectron2

另外，为了快速编译，还可安装ninja，ninja安装方法可参见这篇文章；

三、数据准备

1. 标注工具选择

• 所需要的环境搭建好以后，接下来便是准备数据集。数据集是深度学习的根源，可以说，数据集的质量决定了模型的质量。说到数据集，这里不得不提一下数据集的标注工具，常用的两种分别为：Via和labelme。二者的最大区别为：（1）Via是以网页形式进行标注，而labelme为线下的本地形式进行标注；（2）Via一次标注多张图片后，得到一个.json文件，而labelme一次标注多张图片后，分别得到多个与图片相对应的.json文件；至于使用哪一个标注工具，因人而异，个人更倾向于使用后者进行数据标注。值得注意的是，数据集的标注是个大工程，标注质量的好坏直接影响到模型的精准度。
• 我们在起初时候使用的是Via进行数据标注，由于Via是网页形式进行标注，收到诸多因素限制（如：网络，标签定义与修改等）我们最终选择了Labelme，并将Via标注得到的.json文件转化为了Labelme格式的.json文件。代码如下：

import io
from io import BytesIO
import binascii
import json
import numpy as np
import PIL
import os.path as osp
from labelme.logger import logger
from labelme import PY2
from labelme import QT4
import PIL.Image
import base64
from labelme import utils
import os
import sys

module_path = os.path.abspath(os.path.join('..'))
if module_path not in sys.path:
    sys.path.append(module_path)
# 定义label标签类别
from networkx.generators.tests.test_small import null

def RecongnitionRet (ret ):
	# 这里的“****”需要更改成为你所需要的类别名称，如：person、ballon等，
    if ret == 1:
        label_ret = "****"
    if ret == 2:
        label_ret = "****"
    if ret == 3:
        label_ret = "****"

    return label_ret

# 生成imagedata
def load_image_file(filename):
    try:
        image_pil = PIL.Image.open(filename)
    except IOError:
        logger.error('Failed opening image file: {}'.format(filename))
        return

    # 指定图片格式
    image_pil = utils.apply_exif_orientation(image_pil)
    image = np.array(image_pil)
    with io.BytesIO() as f:
        ext = osp.splitext(filename)[1].lower()
        if PY2 and QT4:
            format = 'PNG'
        elif ext in ['.jpg', '.jpeg']:
            format = 'JPEG'
        else:
            format = 'PNG'
        image_pil.save(f, format=format)
        f.seek(0)
        return f.read(), image

# 转化成json字典
def dict_json(flags, imageData, shapes, imagePath, imageHeight=100, imageWidth=100):
    '''
    :param imageData: str
    :param shapes: list
    :param imagePath: str
    :return: dict
    '''
    return {"version": "4.5.6", "flags": flags, "shapes": shapes,
            'imagePath': imagePath.split('/')[-1], "imageData": imageData, 'imageHeight': imageHeight,
            'imageWidth': imageWidth}

if __name__ == '__main__':
    # print("hello word ... ")
    # 读取json文件内容,返回字典格式，此处需要更改文件路径及json名称
    with open('/home/ytf/PycharmProjects/tools/via_project_5Jan2021_17h40m_json.json', 'r',encoding='utf8') as fp:
        json_data = json.load(fp)
    key_list = json_data.keys()

    # 遍历via_json中所有标注的图片
    for k in key_list:
        print("keynum =", k)
        # 获取json文本内容
        list_data = json_data[k]
        print("list_data ===", list_data)
        # 获取图片名字
        filename = list_data.get('filename')
        imagePath = filename
        # 生成图片数据, 此处需要更改文件路径及名称
        image_Path = '/home/ytf/PycharmProjects/tools/' + imagePath
        imageData, image = load_image_file(image_Path)
        imageData = base64.b64encode(imageData).decode('utf-8')
        imageHeight = image.shape[0]
        imageWidth = image.shape[1]
        # 获取标注框轮廓区域
        regions = list_data.get('regions')
        # 获取轮廓属性
        newshapes = []
        shapes = []
        for i in range(0, len(regions)):
            regions_ret = regions[i]
            # 轮廓属性
            shape_attributes = regions_ret.get('shape_attributes')
            # 轮廓属性-类型
            shape_type = shape_attributes.get('name')
            # 轮廓属性-所有横坐标点
            all_points_x = shape_attributes.get('all_points_x')
            # 轮廓属性-所有纵坐标点
            all_points_y = shape_attributes.get('all_points_y')
            # 轮廓属性-标签
            region_attributes = regions_ret.get('region_attributes')
            reg_val = region_attributes.get('shapes')
            label = RecongnitionRet (int(reg_val ))
            # 横纵坐标组合
            points = []
            x = []
            for m in range(0, len(all_points_x)):
                point = [np.float(all_points_x[m]), np.float(all_points_y[m])]
                points.append(point)
            # 组合新的json成员
            flags = {}
            line_color = None
            fill_color = None
            group_id = None
            newshapes.append(
                {"label": label,
                "points": points,
                "group_id": group_id,
                "shape_type": shape_type,
                "flags": flags})
            # print("newshapes ==", newshapes)
        data_final = dict_json(flags, imageData, newshapes,imagePath, imageHeight, 	
        			imageWidth)
        json_file = imagePath[:-4] + '.json'
        print("json_file =", json_file)
        json.dump(data_final, open(json_file, 'w'))

• Labelme=4.5.6下载链接：

  下载链接

2. 加载数据集合

# 加载数据集合
    register_coco_instances("dataset", {}, "../datasets/project/train_coco/test_coco.json", ".")

上代码中的test_coco.json文件是将标注好的数据，通过调用labelme2coco.py文件生成，具体实现代码如下：

import labelme2coco

# set directory that contains labelme annotations and image files
labelme_folder = "../datasets/project/train/"		# 图片与.json文件路径
# set path for coco json to be saved
save_json_path = "../datasets/project/train_coco/test_coco.json" # 保存路径
# convert labelme annotations to coco
labelme2coco.convert(labelme_folder, save_json_path)

四、参数调整

• 数据集准备好以后，接下来，便可进行训练参数的调整。Detectron2工程中集有基于以ResNet50、ResNet101为骨干网络的目标检测与分割的预训练权重文件，将相应的参数进行调整便可实现预期效果。不可谓不强大呀。
  实际使用中，参数可进行调整，主要需要调整的有：

# 根据实际的使用需要，可在mask_rcnn_R_50_FPN_3x.yaml等文件中，配置需要更改的参数
MODEL:
# 选用哪种骨干网络，主要有R-50, R-101两种；
  WEIGHTS: "detectron2://ImageNetPretrained/MSRA/R-50.pkl"
  MASK_ON: True
  RESNETS:
    DEPTH: 50
  
  STEPS: (210000, 250000)  # 在这个区间内，学习率进行周期性的降低
  MAX_ITER: 270000		# 最大迭代次数
  IMS_PER_BATCH: 8		# 批大小，每次迭代加载的图片数
  BASE_LR: 0.002		# 学习率
INPUT:
  MIN_SIZE_TRAIN: (600, 700, 800, 900)  # 最小训练尺寸，这里是多尺度训练，可提高模型的精准度，但		
  										# 会加大模型训练耗时
  MAX_SIZE_TRAIN: 1333                  # 训练的最大尺寸

• 更多参数说明可看Detectron2官方文档，链接如下：

  Detectron2官方文档接描述

五、模型训练

• 首先，需要加载上文中生成的test_coco.json文件与及参数文件，若该文件加载错误，则训练无法进行。

# 模型训练
def train():
	# 加载数据
    register_coco_instances("dataset", {}, "../datasets/project/train_coco/test_coco.json", ".")
	# 配置参数
    cfg = get_cfg()                            # 获取配置参数
    cfg.OUTPUT_DIR = "../projectName-output"   # 模型输出路径
    cfg.merge_from_file("../configs/project/mask_rcnn_R_50_FPN_3x.yaml") # 选用resnet50作		为卷积网络                   
    cfg.DATASETS.TRAIN = ("dataset",)
    cfg.DATASETS.TEST = ()  # no metrics implemented for this dataset
    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")  # initialize from model zoo

• 数据增广：为了提高模型的精准度及泛化能力，可对数据集进行几何变换和光学变换，即数据增广。光学变换主要包括亮度亮度和对比度等随机调整，可调整图片像素值大小，而并不会改变图像尺寸；集合变换主要包括剪裁和镜像等操作，主要为尺度上的变换。在Detectron2中提供了十多种数据增强的接口，可根据具体需要进行调用。

六、模型测试

• 得到训练模型后，便可对模型进行测试。

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("dataset", {}, "../datasets/"
                                                      "project/"
                                                      "train_coco/"
                                                      "test_coco_" +
                                                      ".json", ".")
    metadata = MetadataCatalog.get("dataset")
    cfg.merge_from_file("../configs/project/mask_rcnn_R_50_FPN_3x.yaml")
    cfg.DATASETS.TRAIN = ("dataset",)
    cfg.DATASETS.TEST = ()

    cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_" + use_model + ".pth")  # path to the model we just trained
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3 # 类别数量，此参数若是不能与训练时的保持一致，测试结果将会严重错误（如：类别错位，只显示一种类别等）
    score_test = 0.5
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = score_test  # 低于50分的box过滤掉
    dataset_dicts = DatasetCatalog.get("dataset")
    # 测试图片保存路径
    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...")

Detectron2默认一张图片上最大可检测目标的数量为100，如果需要检测的目标数量超过了100， 那么就需要对该参数进行更改；

TEST:
  DETECTIONS_PER_IMAGE: 100			# 可检测目标的最大数量

七、总结

• Detectron2做为当下目标检测的神器之一，其功能强大不容置喙。本文只是简单的介绍了下如何利用其进行训练自己的数据集。后续将会继续深入研究其底层原理，主要包括：骨干网络ResNet50、ResNet101的卷积原理，所使用的激活函数，损失函数，优化器等设计与实现方法。

八、参考

1. https://blog.csdn.net/qiuguolu1108/article/details/103842556
2. https://zhuanlan.zhihu.com/p/96931265