更换文档检测模型

paddle_detection/tools/eval_mot.py

@@ -0,0 +1,144 @@
+# 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
+
+import os
+import sys
+
+# add python path of PaddleDetection to sys.path
+parent_path = os.path.abspath(os.path.join(__file__, *(['..'] * 2)))
+sys.path.insert(0, parent_path)
+
+# ignore warning log
+import warnings
+warnings.filterwarnings('ignore')
+
+import paddle
+
+from ppdet.core.workspace import load_config, merge_config
+from ppdet.utils.check import check_gpu, check_npu, check_xpu, check_mlu, check_version, check_config
+from ppdet.utils.cli import ArgsParser
+from ppdet.engine import Tracker
+
+
+def parse_args():
+    parser = ArgsParser()
+    parser.add_argument(
+        "--det_results_dir",
+        type=str,
+        default='',
+        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).')
+    parser.add_argument(
+        '--scaled',
+        type=bool,
+        default=False,
+        help="Whether coords after detector outputs are scaled, False in JDE YOLOv3 "
+        "True in general detector.")
+    args = parser.parse_args()
+    return args
+
+
+def run(FLAGS, cfg):
+    dataset_dir = cfg['EvalMOTDataset'].dataset_dir
+    data_root = cfg['EvalMOTDataset'].data_root
+    data_root = '{}/{}'.format(dataset_dir, data_root)
+    seqs = os.listdir(data_root)
+    seqs.sort()
+
+    # build Tracker
+    tracker = Tracker(cfg, mode='eval')
+
+    # load weights
+    if cfg.architecture in ['DeepSORT', 'ByteTrack']:
+        tracker.load_weights_sde(cfg.det_weights, cfg.reid_weights)
+    else:
+        tracker.load_weights_jde(cfg.weights)
+
+    # inference
+    tracker.mot_evaluate(
+        data_root=data_root,
+        seqs=seqs,
+        data_type=cfg.metric.lower(),
+        model_type=cfg.architecture,
+        output_dir=FLAGS.output_dir,
+        save_images=FLAGS.save_images,
+        save_videos=FLAGS.save_videos,
+        show_image=FLAGS.show_image,
+        scaled=FLAGS.scaled,
+        det_results_dir=FLAGS.det_results_dir)
+
+
+def main():
+    FLAGS = parse_args()
+    cfg = load_config(FLAGS.config)
+    merge_config(FLAGS.opt)
+
+    # disable npu in config by default
+    if 'use_npu' not in cfg:
+        cfg.use_npu = False
+
+    # disable xpu in config by default
+    if 'use_xpu' not in cfg:
+        cfg.use_xpu = False
+
+    if 'use_gpu' not in cfg:
+        cfg.use_gpu = False
+
+    # disable mlu in config by default
+    if 'use_mlu' not in cfg:
+        cfg.use_mlu = False
+
+    if cfg.use_gpu:
+        place = paddle.set_device('gpu')
+    elif cfg.use_npu:
+        place = paddle.set_device('npu')
+    elif cfg.use_xpu:
+        place = paddle.set_device('xpu')
+    elif cfg.use_mlu:
+        place = paddle.set_device('mlu')
+    else:
+        place = paddle.set_device('cpu')
+
+    check_config(cfg)
+    check_gpu(cfg.use_gpu)
+    check_npu(cfg.use_npu)
+    check_xpu(cfg.use_xpu)
+    check_mlu(cfg.use_mlu)
+    check_version()
+
+    run(FLAGS, cfg)
+
+
+if __name__ == '__main__':
+    main()