infer.py 9.4 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18
# Copyright (c) 2019 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

Q
qingqing01 已提交
19
import os, sys
20 21 22 23
# add python path of PadleDetection to sys.path
parent_path = os.path.abspath(os.path.join(__file__, *(['..'] * 2)))
if parent_path not in sys.path:
    sys.path.append(parent_path)
24

Q
qingqing01 已提交
25
import glob
26
import numpy as np
27
import six
28 29
from PIL import Image

30
import paddle
31 32 33 34 35 36
from paddle import fluid

from ppdet.core.workspace import load_config, merge_config, create

from ppdet.utils.eval_utils import parse_fetches
from ppdet.utils.cli import ArgsParser
37
from ppdet.utils.check import check_gpu, check_version, check_config, enable_static_mode
38 39 40
from ppdet.utils.visualizer import visualize_results
import ppdet.utils.checkpoint as checkpoint

41 42
from ppdet.data.reader import create_reader

43 44 45 46 47 48 49 50 51 52 53 54
import logging
FORMAT = '%(asctime)s-%(levelname)s: %(message)s'
logging.basicConfig(level=logging.INFO, format=FORMAT)
logger = logging.getLogger(__name__)


def get_save_image_name(output_dir, image_path):
    """
    Get save image name from source image path.
    """
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
K
Kaipeng Deng 已提交
55
    image_name = os.path.split(image_path)[-1]
56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72
    name, ext = os.path.splitext(image_name)
    return os.path.join(output_dir, "{}".format(name)) + ext


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):
W
wangguanzhong 已提交
73
        return [infer_img]
74

W
wangguanzhong 已提交
75
    images = set()
76 77 78 79 80 81
    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:
W
wangguanzhong 已提交
82 83
        images.update(glob.glob('{}/*.{}'.format(infer_dir, ext)))
    images = list(images)
84 85 86 87 88 89 90 91 92 93 94

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

    return images


def main():
    cfg = load_config(FLAGS.config)

    merge_config(FLAGS.opt)
95
    check_config(cfg)
96 97
    # check if set use_gpu=True in paddlepaddle cpu version
    check_gpu(cfg.use_gpu)
W
wangguanzhong 已提交
98 99
    # check if paddlepaddle version is satisfied
    check_version()
100

101 102
    main_arch = cfg.architecture

103
    dataset = cfg.TestReader['dataset']
104 105

    test_images = get_test_images(FLAGS.infer_dir, FLAGS.infer_img)
106
    dataset.set_images(test_images)
107 108 109 110 111 112 113 114 115 116

    place = fluid.CUDAPlace(0) if cfg.use_gpu else fluid.CPUPlace()
    exe = fluid.Executor(place)

    model = create(main_arch)

    startup_prog = fluid.Program()
    infer_prog = fluid.Program()
    with fluid.program_guard(infer_prog, startup_prog):
        with fluid.unique_name.guard():
117 118 119
            inputs_def = cfg['TestReader']['inputs_def']
            inputs_def['iterable'] = True
            feed_vars, loader = model.build_inputs(**inputs_def)
120 121 122
            test_fetches = model.test(feed_vars)
    infer_prog = infer_prog.clone(True)

123
    reader = create_reader(cfg.TestReader, devices_num=1)
W
wangguanzhong 已提交
124
    loader.set_sample_list_generator(reader, place)
125 126 127

    exe.run(startup_prog)
    if cfg.weights:
128
        checkpoint.load_params(exe, infer_prog, cfg.weights)
129 130

    # parse infer fetches
131
    assert cfg.metric in ['COCO', 'VOC', 'OID', 'WIDERFACE'], \
132
            "unknown metric type {}".format(cfg.metric)
133
    extra_keys = []
134
    if cfg['metric'] in ['COCO', 'OID']:
135
        extra_keys = ['im_info', 'im_id', 'im_shape']
136
    if cfg['metric'] == 'VOC' or cfg['metric'] == 'WIDERFACE':
137
        extra_keys = ['im_id', 'im_shape']
138 139 140 141
    keys, values, _ = parse_fetches(test_fetches, infer_prog, extra_keys)

    # parse dataset category
    if cfg.metric == 'COCO':
G
Guanghua Yu 已提交
142
        from ppdet.utils.coco_eval import bbox2out, mask2out, segm2out, get_category_info
143 144
    if cfg.metric == 'OID':
        from ppdet.utils.oid_eval import bbox2out, get_category_info
145 146
    if cfg.metric == "VOC":
        from ppdet.utils.voc_eval import bbox2out, get_category_info
147
    if cfg.metric == "WIDERFACE":
148
        from ppdet.utils.widerface_eval_utils import bbox2out, lmk2out, get_category_info
149

150 151 152 153
    anno_file = dataset.get_anno()
    with_background = dataset.with_background
    use_default_label = dataset.use_default_label

154 155 156
    clsid2catid, catid2name = get_category_info(anno_file, with_background,
                                                use_default_label)

157 158 159 160 161 162
    # whether output bbox is normalized in model output layer
    is_bbox_normalized = False
    if hasattr(model, 'is_bbox_normalized') and \
            callable(model.is_bbox_normalized):
        is_bbox_normalized = model.is_bbox_normalized()

走神的阿圆's avatar
走神的阿圆 已提交
163 164
    # use VisualDL to log image
    if FLAGS.use_vdl:
165
        assert six.PY3, "VisualDL requires Python >= 3.5"
走神的阿圆's avatar
走神的阿圆 已提交
166 167 168 169
        from visualdl import LogWriter
        vdl_writer = LogWriter(FLAGS.vdl_log_dir)
        vdl_image_step = 0
        vdl_image_frame = 0  # each frame can display ten pictures at most.
170

171
    imid2path = dataset.get_imid2path()
W
wangguanzhong 已提交
172
    for iter_id, data in enumerate(loader()):
173
        outs = exe.run(infer_prog,
W
wangguanzhong 已提交
174
                       feed=data,
175 176 177 178 179 180 181
                       fetch_list=values,
                       return_numpy=False)
        res = {
            k: (np.array(v), v.recursive_sequence_lengths())
            for k, v in zip(keys, outs)
        }
        logger.info('Infer iter {}'.format(iter_id))
W
wangguanzhong 已提交
182 183
        if 'TTFNet' in cfg.architecture:
            res['bbox'][1].append([len(res['bbox'][0])])
184 185 186 187
        if 'CornerNet' in cfg.architecture:
            from ppdet.utils.post_process import corner_post_process
            post_config = getattr(cfg, 'PostProcess', None)
            corner_post_process(res, post_config, cfg.num_classes)
188 189 190

        bbox_results = None
        mask_results = None
G
Guanghua Yu 已提交
191
        segm_results = None
192
        lmk_results = None
193 194 195 196 197
        if 'bbox' in res:
            bbox_results = bbox2out([res], clsid2catid, is_bbox_normalized)
        if 'mask' in res:
            mask_results = mask2out([res], clsid2catid,
                                    model.mask_head.resolution)
G
Guanghua Yu 已提交
198 199
        if 'segm' in res:
            segm_results = segm2out([res], clsid2catid)
200 201
        if 'landmark' in res:
            lmk_results = lmk2out([res], is_bbox_normalized)
202 203 204 205 206 207

        # visualize result
        im_ids = res['im_id'][0]
        for im_id in im_ids:
            image_path = imid2path[int(im_id)]
            image = Image.open(image_path).convert('RGB')
208

走神的阿圆's avatar
走神的阿圆 已提交
209 210
            # use VisualDL to log original image
            if FLAGS.use_vdl:
211
                original_image_np = np.array(image)
走神的阿圆's avatar
走神的阿圆 已提交
212 213 214
                vdl_writer.add_image(
                    "original/frame_{}".format(vdl_image_frame),
                    original_image_np, vdl_image_step)
215

216
            image = visualize_results(image,
J
jerrywgz 已提交
217 218
                                      int(im_id), catid2name,
                                      FLAGS.draw_threshold, bbox_results,
G
Guanghua Yu 已提交
219
                                      mask_results, segm_results, lmk_results)
W
wangguanzhong 已提交
220

走神的阿圆's avatar
走神的阿圆 已提交
221 222
            # use VisualDL to log image with bbox
            if FLAGS.use_vdl:
223
                infer_image_np = np.array(image)
走神的阿圆's avatar
走神的阿圆 已提交
224 225 226 227 228 229
                vdl_writer.add_image("bbox/frame_{}".format(vdl_image_frame),
                                     infer_image_np, vdl_image_step)
                vdl_image_step += 1
                if vdl_image_step % 10 == 0:
                    vdl_image_step = 0
                    vdl_image_frame += 1
230

231 232
            save_name = get_save_image_name(FLAGS.output_dir, image_path)
            logger.info("Detection bbox results save in {}".format(save_name))
J
jerrywgz 已提交
233
            image.save(save_name, quality=95)
234 235 236


if __name__ == '__main__':
237
    enable_static_mode()
238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253
    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.")
J
jerrywgz 已提交
254 255 256 257 258
    parser.add_argument(
        "--draw_threshold",
        type=float,
        default=0.5,
        help="Threshold to reserve the result for visualization.")
259
    parser.add_argument(
走神的阿圆's avatar
走神的阿圆 已提交
260
        "--use_vdl",
261 262
        type=bool,
        default=False,
走神的阿圆's avatar
走神的阿圆 已提交
263
        help="whether to record the data to VisualDL.")
264
    parser.add_argument(
走神的阿圆's avatar
走神的阿圆 已提交
265
        '--vdl_log_dir',
266
        type=str,
走神的阿圆's avatar
走神的阿圆 已提交
267 268
        default="vdl_log_dir/image",
        help='VisualDL logging directory for image.')
269
    FLAGS = parser.parse_args()
W
wangguanzhong 已提交
270
    main()