infer.py 6.6 KB
Newer Older
Q
qingqing01 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
M
Manuel Garcia 已提交
18 19 20 21

import os
import sys

Q
qingqing01 已提交
22 23
# add python path of PadleDetection to sys.path
parent_path = os.path.abspath(os.path.join(__file__, *(['..'] * 2)))
24
sys.path.insert(0, parent_path)
Q
qingqing01 已提交
25

G
Guanghua Yu 已提交
26 27 28
# ignore warning log
import warnings
warnings.filterwarnings('ignore')
Q
qingqing01 已提交
29
import glob
W
wangxinxin08 已提交
30
import ast
Q
qingqing01 已提交
31 32

import paddle
K
Kaipeng Deng 已提交
33 34
from ppdet.core.workspace import load_config, merge_config
from ppdet.engine import Trainer
H
houj04 已提交
35
from ppdet.utils.check import check_gpu, check_npu, check_xpu, check_version, check_config
W
wangxinxin08 已提交
36
from ppdet.utils.cli import ArgsParser, merge_args
37
from ppdet.slim import build_slim_model
Q
qingqing01 已提交
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

from ppdet.utils.logger import setup_logger
logger = setup_logger('train')


def parse_args():
    parser = ArgsParser()
    parser.add_argument(
        "--infer_dir",
        type=str,
        default=None,
        help="Directory for images to perform inference on.")
    parser.add_argument(
        "--infer_img",
        type=str,
        default=None,
        help="Image path, has higher priority over --infer_dir")
    parser.add_argument(
        "--output_dir",
        type=str,
        default="output",
        help="Directory for storing the output visualization files.")
    parser.add_argument(
        "--draw_threshold",
        type=float,
        default=0.5,
        help="Threshold to reserve the result for visualization.")
65 66 67 68 69
    parser.add_argument(
        "--slim_config",
        default=None,
        type=str,
        help="Configuration file of slim method.")
Q
qingqing01 已提交
70 71 72 73
    parser.add_argument(
        "--use_vdl",
        type=bool,
        default=False,
74
        help="Whether to record the data to VisualDL.")
Q
qingqing01 已提交
75 76 77 78 79
    parser.add_argument(
        '--vdl_log_dir',
        type=str,
        default="vdl_log_dir/image",
        help='VisualDL logging directory for image.')
C
cnn 已提交
80
    parser.add_argument(
W
Wenyu 已提交
81
        "--save_results",
C
cnn 已提交
82 83
        type=bool,
        default=False,
W
Wenyu 已提交
84
        help="Whether to save inference results to output_dir.")
85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115
    parser.add_argument(
        "--slice_infer",
        action='store_true',
        help="Whether to slice the image and merge the inference results for small object detection."
    )
    parser.add_argument(
        '--slice_size',
        nargs='+',
        type=int,
        default=[640, 640],
        help="Height of the sliced image.")
    parser.add_argument(
        "--overlap_ratio",
        nargs='+',
        type=float,
        default=[0.25, 0.25],
        help="Overlap height ratio of the sliced image.")
    parser.add_argument(
        "--combine_method",
        type=str,
        default='nms',
        help="Combine method of the sliced images' detection results, choose in ['nms', 'nmm', 'concat']."
    )
    parser.add_argument(
        "--match_threshold",
        type=float,
        default=0.6,
        help="Combine method matching threshold.")
    parser.add_argument(
        "--match_metric",
        type=str,
F
Feng Ni 已提交
116
        default='ios',
117
        help="Combine method matching metric, choose in ['iou', 'ios'].")
W
wangxinxin08 已提交
118 119 120 121 122
    parser.add_argument(
        "--visualize",
        type=ast.literal_eval,
        default=True,
        help="Whether to save visualize results to output_dir.")
Q
qingqing01 已提交
123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157
    args = parser.parse_args()
    return args


def get_test_images(infer_dir, infer_img):
    """
    Get image path list in TEST mode
    """
    assert infer_img is not None or infer_dir is not None, \
        "--infer_img or --infer_dir should be set"
    assert infer_img is None or os.path.isfile(infer_img), \
            "{} is not a file".format(infer_img)
    assert infer_dir is None or os.path.isdir(infer_dir), \
            "{} is not a directory".format(infer_dir)

    # infer_img has a higher priority
    if infer_img and os.path.isfile(infer_img):
        return [infer_img]

    images = set()
    infer_dir = os.path.abspath(infer_dir)
    assert os.path.isdir(infer_dir), \
        "infer_dir {} is not a directory".format(infer_dir)
    exts = ['jpg', 'jpeg', 'png', 'bmp']
    exts += [ext.upper() for ext in exts]
    for ext in exts:
        images.update(glob.glob('{}/*.{}'.format(infer_dir, ext)))
    images = list(images)

    assert len(images) > 0, "no image found in {}".format(infer_dir)
    logger.info("Found {} inference images in total.".format(len(images)))

    return images


K
Kaipeng Deng 已提交
158 159 160 161 162
def run(FLAGS, cfg):
    # build trainer
    trainer = Trainer(cfg, mode='test')

    # load weights
K
Kaipeng Deng 已提交
163
    trainer.load_weights(cfg.weights)
K
Kaipeng Deng 已提交
164 165 166 167 168

    # get inference images
    images = get_test_images(FLAGS.infer_dir, FLAGS.infer_img)

    # inference
169 170 171 172 173 174 175 176 177 178
    if FLAGS.slice_infer:
        trainer.slice_predict(
            images,
            slice_size=FLAGS.slice_size,
            overlap_ratio=FLAGS.overlap_ratio,
            combine_method=FLAGS.combine_method,
            match_threshold=FLAGS.match_threshold,
            match_metric=FLAGS.match_metric,
            draw_threshold=FLAGS.draw_threshold,
            output_dir=FLAGS.output_dir,
W
wangxinxin08 已提交
179 180
            save_results=FLAGS.save_results,
            visualize=FLAGS.visualize)
181 182 183 184 185
    else:
        trainer.predict(
            images,
            draw_threshold=FLAGS.draw_threshold,
            output_dir=FLAGS.output_dir,
W
wangxinxin08 已提交
186 187
            save_results=FLAGS.save_results,
            visualize=FLAGS.visualize)
Q
qingqing01 已提交
188 189 190 191 192


def main():
    FLAGS = parse_args()
    cfg = load_config(FLAGS.config)
W
wangxinxin08 已提交
193
    merge_args(cfg, FLAGS)
Q
qingqing01 已提交
194
    merge_config(FLAGS.opt)
195

196 197 198 199
    # disable npu in config by default
    if 'use_npu' not in cfg:
        cfg.use_npu = False

H
houj04 已提交
200 201 202 203
    # disable xpu in config by default
    if 'use_xpu' not in cfg:
        cfg.use_xpu = False

204 205 206 207
    if cfg.use_gpu:
        place = paddle.set_device('gpu')
    elif cfg.use_npu:
        place = paddle.set_device('npu')
H
houj04 已提交
208 209
    elif cfg.use_xpu:
        place = paddle.set_device('xpu')
210 211
    else:
        place = paddle.set_device('cpu')
212

213
    if FLAGS.slim_config:
214 215
        cfg = build_slim_model(cfg, FLAGS.slim_config, mode='test')

Q
qingqing01 已提交
216 217
    check_config(cfg)
    check_gpu(cfg.use_gpu)
218
    check_npu(cfg.use_npu)
H
houj04 已提交
219
    check_xpu(cfg.use_xpu)
220
    check_version('develop')
Q
qingqing01 已提交
221

K
Kaipeng Deng 已提交
222
    run(FLAGS, cfg)
Q
qingqing01 已提交
223 224 225 226


if __name__ == '__main__':
    main()