Detectron 源码解析-利用预训练模型检测自定义图片

发表于   |   更新于 | 分类于 | 评论数： | 阅读次数：
本文字数： 9.5k | 阅读时长 ≈ 9 分钟

源码文件

  • tools/infer_simple.py

运行指令

Detectron 是一个非常庞杂的代码项目, 因此, 在对 Detecron 内部的主要模型进行解析之前, 我们首先要知道如何使用已经训练好的模型. 也就是先要知道如何利用完善的 Detectron 目标检测模型框架. 根据官方的使用教程, 要利用 Detectron 检测用户提供的自定义图片, 需要在 shell 中运行下面的指令:

1
2
3
4
5
6
python tools/infer_simple.py \
    --cfg configs/12_2017_baselines/e2e_mask_rcnn_R-101-FPN_2x.yaml \
    --output-dir /tmp/detectron-visualizations \
    --image-ext igp \
    --wts https://s3-us-west-2.amazonaws.com/detectron/35861858/12_2017_baselines/e2e_mask_rcnn_R-101-FPN_2x.yaml.02_32_51.SgT4y1cO/output/train/coco_2014_train:coco_2014_valminusminival/generalized_rcnn/model_final.pkl \
    demo

在上面的指令中, Detectron 会自动下载 --wts 参数指定了模型的下载链接, 最终下载的模型会根据 Detectron/detectron/core/config.py 中的 __C.DOWNLOAD_CACHE 参数决定, 该参数的默认值为 /tmp/detectron-download-cache, 可以根据你自己的需要进行修改, --output-dir 参数指定了返回的结果存放的位置, 推荐你修改成自己期望的路径, 结果在默认情况下会以 PDF 的格式返回. 最后的demo 指示的是待检测图谱存放的路径, 位于Detectron/demo/文件夹下, 这里可以直接指定的原因是因为在代码中使用了下面的方式定义此参数

1
2
3
parser.add_argument(
    'im_or_folder', help='image or folder of images', default=None
) # 注意第一个参数定义为'im_or_folder', 而不是'--im_or_folder'

tools/infer_simple 文件总览

接下来, 我们根据这个脚本文件, 逐步解析 Detectron 的工作原理. 为了更加有条理的解析, 我们会根据文件中的代码功能, 将文件按照模块进行解析, 主要包含以下几个部分:

  • 导入的包及函数: 介绍导入的函数和变量的含义及作用
  • 命令行参数: 介绍该文件使用的命令行参数的用途及注意事项
  • 主程序: 最重要的一部分, 可能包含大量其他文件中的函数代码, 会重点讲解该文件.

文件代码的概览如下所示:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
# ./tools/infer_simple

# 导入的包及函数
from __future__ import absolute_import
#...
from caffe2.python import workspace
#...
from detectron.utils.logging import setup_logging
#...

# 命令行参数
def parse_args():
    parser = argparse.ArgumentParser(description="End-to-end inference")
    parser.add_argument(...)
    #...
    return parser.parse_args()

# 主程序
def main(args):
    logger = logging.getLogger(__name__)
    merge_cfg_from_file(args.cfg)
    #...

if __name__ == "__main__":
    workspace.GlobalInit(["caffe2", "--caffe2_log_level=0"])
    setup_logging(__name__) # 调用 detectron/utils/logging.py中的函数
    args = parse_args()
    main(args)

导入的包及函数

首先看一下该文件导入了那些包和函数, 以及它们的作用:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
# ./tools/infer_simple

# python2兼容包
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

# 标准常用包, 下面的包为常用的一些包, 不多做介绍
from collections import defaultdict
import argparse
import cv2
import glob
import logging
import os
import sys
import time

from caffe2.python import workspace

# assert_and_infer_cfg函数会将所有的 cfg 选项参数设置为只读, 防止修改
from detectron.core.config import assert_and_infer_cfg
# 从config.py文件中导入了cfg选项, 该选项包含了所有 Detectron 模型的选择参数, 因此体量非常大(约1000行代码)
from detectron.core.config import cfg
# merge_cfg_from_file函数用于将一个 yaml config 文件中的 cfg 选项融合到全局 cfg 当中(即core/config.py中的cfg)
from detectron.core.config import merge_cfg_from_file

#
from detectron.utils.io import cache_url

#
from detectron.utils.logging import setup_logging

#
from detectron.utils.timer import Timer

#
import detectron.core.test_engine as infer_engine

#
import detectron.datasets.dummy_datasets as dummy_datasets

#
import detectron.utils.c2 as c2_utils

#
import detectron.utils.vis as vis_utils

在上面的代码中, 首先导入了 cfg 相关的变量和函数, 它们大致作用已经在代码中的注释中介绍, 如果你还想深入了解这些函数的具体实现, 那么可以看关于全局配置选项及相关函数. 另外还有用于设置输出日志信息的函数 setup_logging, 详细解析可以看日志输出控制及训练状态跟踪.

命令行参数

下面的代码定义了该文件可能会用到的命令行参数, 其中有些参数你已经在文章开始的运行指令中见到过, 下面我们就来看一下这些参数的定义及功能.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
# Detectron/tools/infer_simple.py

def parse_args():
    parser = argparse.ArgumentParser(description='End-to-end inference')
    parser.add_argument(
        '--cfg',
        dest='cfg',
        help='cfg model file (/path/to/model_config.yaml)',
        default=None,
        type=str
    )
    parser.add_argument(
        '--wts',
        dest='weights',
        help='weights model file (/path/to/model_weights.pkl)',
        default=None,
        type=str
    )
    parser.add_argument(
        '--output-dir',
        dest='output_dir',
        help='directory for visualization pdfs (default: /tmp/infer_simple)',
        default='/tmp/infer_simple',
        type=str
    )
    parser.add_argument(
        '--image-ext',
        dest='image_ext',
        help='image file name extension (default: jpg)',
        default='jpg',
        type=str
    )
    parser.add_argument(
        '--always-out',
        dest='out_when_no_box',
        help='output image even when no object is found',
        action='store_true'
    )
    parser.add_argument(
        '--output-ext',
        dest='output_ext',
        help='output image file format (default: pdf)',
        default='pdf',
        type=str
    )
    parser.add_argument(
        '--thresh',
        dest='thresh',
        help='Threshold for visualizing detections',
        default=0.7,
        type=float
    )
    parser.add_argument(
        '--kp-thresh',
        dest='kp_thresh',
        help='Threshold for visualizing keypoints',
        default=2.0,
        type=float
    )
    parser.add_argument(
        'im_or_folder', help='image or folder of images', default=None
    )# 注意这里使用的是'im_or_folder', 而不是'--im_or_folder'
    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(1)
    return parser.parse_args()

主程序

mian() 函数为该脚本文件的主要函数, 其调用了多个其他文件的函数及参数, 主要有以下几个:

  • merge_cfg_from_file(): 将 yaml config 文件中的配置信息融合到 detectron/core/config.py 中的全局配置信息中.
  • cache_url(): 下载url指定的文件到cache_dir当中, 同时返回缓存文件的路径, 如果缓存文件直接存在, 则可以直接返回. 如果传入的第一个参数不是url类型的, 也直接将其返回.
  • assert_and_infer_cfg(): 当完成所有的 cfg 参数设置以后, 调用此函数, 默认情况下, 该函数会将全局的 cfg 标记为只读类型(immutable), 以此来保护在脚本执行过程中 cfg 选项的值不被修改.(如果允许修改, 则在debug的时候会很难)
  • infer_engine.initialize_model_from_cfg(): 该函数会根据全局配置信息(global cfg)创建模型并对其进行初始化. 在将 caffe2 的工作空间中创建网络并加载 test-time 权重. 关于此函数的详细说明可以查看模型推演源码解析
  • dummy_datasets.get_coco_dataset(): 该函数用于加载 coco 数据集, 关于这部分的详细解析可以查看数据载入详细解析
  • vis_utils.vis_one_image(): 该函数用于将模型输出的数值结果以可视化的方式呈现, 关于这部分的详细解析可以查看Detectron 结果可视化
  • c2_utils.NamedCudaScopr(): 提供 GPU name scope 和 CUDA device scope 功能, 需要结合 with 使用, 位于detectron/utils/c2.py文件, 详细解析请看工具函数完整解读

代码解析如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
def main(args):
    logger = logging.getLogger(__name__)
    # 将指定的 yaml config 的选项融合到全局 cfg 当中去
    merge_cfg_from_file(args.cfg)
    cfg.NUM_GPUS = 1
    # 该函数用于下载args.weights指定的模型到cfg.DOWNLOAD_CACHE目录中,
    # 如果模型已经存在, 或者args.weights不是url类型, 则直接返回文件路径. 否则, 会在下载完成后返回文件路径
    args.weights = cache_url(args.weights, cfg.DOWNLOAD_CACHE)
    # 将所有的 cfg 选项参数设置为只读, 防止修改, 这样做易于调试
    assert_and_infer_cfg(cache_urls=False)

    # 确保 detectron/core/config.py 文件中的关键参数设置正确
    assert not cfg.MODEL.RPN_ONLY, \
        'RPN models are not supported'
    assert not cfg.TEST.PRECOMPUTED_PROPOSALS, \
        'Models that require precomputed proposals are not supported'

    # 根据cfg信息创建模型, 并下载/加载指定的权重文件
    model = infer_engine.initialize_model_from_cfg(args.weights)
    # 获取coco数据集的类别(classes).
    # 返回一个 AttrDict 类型, 其内部存储着 'classes: value', value 也是一个字典, 该字典为编号与类型组成的键值对
    # 0:'__background__', 1:'person', ...
    dummy_coco_dataset = dummy_datasets.get_coco_dataset()

    if os.path.isdir(args.im_or_folder):
        # 如果args.im_or_folder给定的是一个文件夹的话, 则会利用glob模块来抓取其中的.jpg文件,
        # 并返回一个可遍历的 _iglob 对象, 其内部存储的是.jpg文件的文件路径.
        im_list = glob.iglob(args.im_or_folder + '/*.' + args.image_ext)
    else:
        # 如果不是文件夹, 则说明传入的是一张图片的路径, 则直接将其放在列表当中.
        im_list = [args.im_or_folder]

    # 不管是 _iglob 类型还是列表类型, 都可以用同样的方式遍历得到图片路径.
    for i, im_name in enumerate(im_list):
        # 根据指定的参数来决定最终的输出路径及输出文件名, 默认的outpu_ext类型为pdf文件
        out_name = os.path.join(
            args.output_dir, '{}'.format(os.path.basename(im_name) + '.' + args.output_ext)
        )
        # 记录相关日志信息
        logger.info('Processing {} -> {}'.format(im_name, out_name))
        # 用opencv读取图片, 注意格式为 GBR
        im = cv2.imread(im_name)
        timers = defaultdict(Timer)
        # 为程序运行时间计时
        t = time.time()
        # 创建一个 GPU name scope 和一个 CUDA device scope.
        with c2_utils.NamedCudaScope(0):
            # 调用 im_detect_all, 用model来处理im图片, 并返回结果.
            # 可以看到, 结果有3类, 分别对应目标检测, 实例分割, 和关键点检测任务.
            cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all(
                model, im, None, timers=timers
            )
        # 记录相关日志信息, 以便打印
        logger.info('Inference time: {:.3f}s'.format(time.time() - t))
        for k, v in timers.items():
            logger.info(' | {}: {:.3f}s'.format(k, v.average_time))
        # 当检测第一张图片时, 打印一些提示信息.
        if i == 0:
            logger.info(
                ' \ Note: inference on the first image will be slower than the '
                'rest (caches and auto-tuning need to warm up)'
            )
        # 将图片进行可视化
        vis_utils.vis_one_image(
            im[:, :, ::-1],  # BGR -> RGB for visualization
            im_name,
            args.output_dir,
            cls_boxes,
            cls_segms,
            cls_keyps,
            dataset=dummy_coco_dataset,
            box_alpha=0.3,
            show_class=True,
            thresh=args.thresh,
            kp_thresh=args.kp_thresh,
            ext=args.output_ext,
            out_when_no_box=args.out_when_no_box
        )

小结

上面的代码是 Detectron 的一个非常简单的示例代码, 涵盖了模型创建, 模型初始化, 模型测试, 结果可视化等功能. 由于这里的实例代码仅仅是调用了封装好的接口, 因此我们虽然成功的通过预训练好的模型对自定义的图片进行了检测, 但是我们可能还是不太清楚其内部到底是怎么实现的. 实际上, 这个示例仅仅是一个“热身”, 如果你想要更加深入的了解 Detectron 的内部实现原理, 可以查看我的其他解析.

PS: 查看这篇总览, 可以大致了解 Detectron 的源码结构, 根据你自己的需要来选择相应的解析文章.