|
- # 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
- import os, sys
- # 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)
-
- # ignore warning log
- import warnings
- warnings.filterwarnings('ignore')
- import glob
-
- import paddle
- from model.core.workspace import load_config, merge_config
- from model.engine import Trainer
- from utils.utils.check import check_gpu, check_version, check_config
- from utils.utils.cli import ArgsParser
-
- from utils.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.")
- parser.add_argument(
- "--use_vdl",
- type=bool,
- default=False,
- help="Whether to record the data to VisualDL.")
- parser.add_argument(
- '--vdl_log_dir',
- type=str,
- default="vdl_log_dir/image",
- help='VisualDL logging directory for image.')
- parser.add_argument(
- "--save_txt",
- type=bool,
- default=False,
- help="Whether to save inference result in txt.")
- 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
-
-
- def run(FLAGS, cfg):
- # build trainer
- trainer = Trainer(cfg, mode='test')
-
- # load weights
- trainer.load_weights(cfg.weights)
-
- # get inference images
- images = get_test_images(FLAGS.infer_dir, FLAGS.infer_img)
-
- # inference
- trainer.predict(
- images,
- draw_threshold=FLAGS.draw_threshold,
- output_dir=FLAGS.output_dir,
- save_txt=FLAGS.save_txt)
-
-
- def main():
- FLAGS = parse_args()
- cfg = load_config(FLAGS.config)
- cfg['use_vdl'] = FLAGS.use_vdl
- cfg['vdl_log_dir'] = FLAGS.vdl_log_dir
- merge_config(FLAGS.opt)
-
- place = paddle.set_device('gpu' if cfg.use_gpu else 'cpu')
-
- if 'norm_type' in cfg and cfg['norm_type'] == 'sync_bn' and not cfg.use_gpu:
- cfg['norm_type'] = 'bn'
-
- check_config(cfg)
- check_gpu(cfg.use_gpu)
- check_version()
-
- run(FLAGS, cfg)
-
-
- if __name__ == '__main__':
- main()
|