infer_mot.py 3.8 KB
Newer Older
G
George Ni 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
# Copyright (c) 2021 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

G
George Ni 已提交
22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45
# 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 paddle
from paddle.distributed import ParallelEnv
from ppdet.core.workspace import load_config, merge_config
from ppdet.engine import Tracker
from ppdet.utils.check import check_gpu, check_version, check_config
from ppdet.utils.cli import ArgsParser

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


def parse_args():
    parser = ArgsParser()
    parser.add_argument(
        '--video_file', type=str, default=None, help='Video name for tracking.')
G
George Ni 已提交
46 47 48 49 50
    parser.add_argument(
        "--image_dir",
        type=str,
        default=None,
        help="Directory for images to perform inference on.")
G
George Ni 已提交
51 52 53 54 55 56 57 58
    parser.add_argument(
        "--data_type",
        type=str,
        default='mot',
        help='Data type of tracking dataset, should be in ["mot", "kitti"]')
    parser.add_argument(
        "--det_results_dir",
        type=str,
59
        default='',
G
George Ni 已提交
60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77
        help="Directory name for detection results.")
    parser.add_argument(
        '--output_dir',
        type=str,
        default='output',
        help='Directory name for output tracking results.')
    parser.add_argument(
        '--save_images',
        action='store_true',
        help='Save tracking results (image).')
    parser.add_argument(
        '--save_videos',
        action='store_true',
        help='Save tracking results (video).')
    parser.add_argument(
        '--show_image',
        action='store_true',
        help='Show tracking results (image).')
78 79 80 81 82
    parser.add_argument(
        "--draw_threshold",
        type=float,
        default=0.5,
        help="Threshold to reserve the result for visualization.")
G
George Ni 已提交
83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102
    args = parser.parse_args()
    return args


def run(FLAGS, cfg):
    # build Tracker
    tracker = Tracker(cfg, mode='test')

    # load weights
    if cfg.architecture in ['DeepSORT']:
        if cfg.det_weights != 'None':
            tracker.load_weights_sde(cfg.det_weights, cfg.reid_weights)
        else:
            tracker.load_weights_sde(None, cfg.reid_weights)
    else:
        tracker.load_weights_jde(cfg.weights)

    # inference
    tracker.mot_predict(
        video_file=FLAGS.video_file,
G
George Ni 已提交
103
        image_dir=FLAGS.image_dir,
G
George Ni 已提交
104 105 106 107 108 109
        data_type=FLAGS.data_type,
        model_type=cfg.architecture,
        output_dir=FLAGS.output_dir,
        save_images=FLAGS.save_images,
        save_videos=FLAGS.save_videos,
        show_image=FLAGS.show_image,
110 111
        det_results_dir=FLAGS.det_results_dir,
        draw_threshold=FLAGS.draw_threshold)
G
George Ni 已提交
112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129


def main():
    FLAGS = parse_args()
    cfg = load_config(FLAGS.config)
    merge_config(FLAGS.opt)

    check_config(cfg)
    check_gpu(cfg.use_gpu)
    check_version()

    place = 'gpu:{}'.format(ParallelEnv().dev_id) if cfg.use_gpu else 'cpu'
    place = paddle.set_device(place)
    run(FLAGS, cfg)


if __name__ == '__main__':
    main()