1""" SiamRPN benchmark
2Code adapted from https://github.com/STVIR/pysot"""
3from glob import glob
4import argparse
5import os
6from multiprocessing import Pool
7from tqdm import tqdm
8from gluoncv.utils.metrics.tracking import OPEBenchmark
9from gluoncv.data.otb.tracking import OTBTracking as OTBDataset
10
11def parse_args():
12    """ benchmark test."""
13    parser = argparse.ArgumentParser(description='tracking evaluation')
14    parser.add_argument('--tracker-path', '-p', type=str, help='test result path')
15    parser.add_argument('--dataset', '-d', default='OTB2015', type=str, help='dataset name')
16    parser.add_argument('--num', '-n', default=1, type=int, help='number of thread to eval')
17    parser.add_argument('--tracker-prefix', '-t', type=str, help='tracker name')
18    parser.add_argument('--show-video-level', '-s', action='store_true')
19    parser.add_argument('--test-dataset', type=str, help='test_json dataset dir')
20    parser.set_defaults(show_video_level=False)
21    opt = parser.parse_args()
22    return opt
23
24def main():
25    """SiamRPN benchmark.
26    evaluation according to txt of test result.now supports benchmark is Success and Precision
27    Currently only supports test OTB 2015 dataset.
28
29    Parameters
30    ----------
31    tracker_path : str, txt of test result path.
32    tracker_prefix : str, model name.
33    test_dataset : str, Path to test label json.
34    """
35    opt = parse_args()
36    tracker_dir = os.path.join(opt.tracker_path, opt.dataset)
37    trackers = glob(os.path.join(opt.tracker_path,
38                                 opt.dataset,
39                                 opt.tracker_prefix+'*'))
40    trackers = [x.split('/')[-1] for x in trackers]
41    assert len(trackers) > 0
42    opt.num = min(opt.num, len(trackers))
43    dataset = OTBDataset(name=opt.dataset, dataset_root=opt.test_dataset, load_img=False)
44    dataset.set_tracker(tracker_dir, trackers)
45    benchmark = OPEBenchmark(dataset)
46    success_ret = {}
47    with Pool(processes=opt.num) as pool:
48        for ret in tqdm(pool.imap_unordered(benchmark.eval_success, trackers),
49                        desc='eval success', total=len(trackers), ncols=100):
50            success_ret.update(ret)
51    precision_ret = {}
52    with Pool(processes=opt.num) as pool:
53        for ret in tqdm(pool.imap_unordered(benchmark.eval_precision, trackers),
54                        desc='eval precision', total=len(trackers), ncols=100):
55            precision_ret.update(ret)
56    benchmark.show_result(success_ret, precision_ret,
57                          show_video_level=opt.show_video_level)
58
59
60if __name__ == '__main__':
61    main()
62