#!/usr/bin/env python

# --------------------------------------------------------
# Fast R-CNN
# Copyright (c) 2015 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Ross Girshick
# --------------------------------------------------------

"""Test a Fast R-CNN network on an image database."""

import _init_paths
from fast_rcnn.test import test_net
from fast_rcnn.config import cfg, cfg_from_file, cfg_from_list
from datasets.factory import get_imdb
import caffe
import argparse
import pprint
import time, os, sys
import pandas as pd


def splitall(path):
    allparts = []
    while 1:
        parts = os.path.split(path)
        if parts[0] == path:  # sentinel for absolute paths
            allparts.insert(0, parts[0])
            break
        elif parts[1] == path:  # sentinel for relative paths
            allparts.insert(0, parts[1])
            break
        else:
            path = parts[0]
            allparts.insert(0, parts[1])
    return allparts


def parse_args():
    """
    Parse input arguments
    """
    parser = argparse.ArgumentParser(description='Test a Fast R-CNN network pipeline')
    parser.add_argument('--gpu', dest='gpu_id', help='GPU id to use',
                        default=0, type=int, required=True)
    parser.add_argument('--dir', dest='dir',
                        help='Directory of the model files',
                        default="", type=str, required=True)
    parser.add_argument('--models', dest='model_files',
                        help='Text file with names of models',
                        default=None, type=str, required=True)
    parser.add_argument('--prototxt', dest='prototxt',
                        help='prototxt', default=None, type=str, required=True)
    parser.add_argument('--imdb', dest='imdb_name',
                        help='dataset to test',
                        default='ped_test_small', type=str, required=True)
    parser.add_argument('--cfg', dest='cfg_file',
                        help='cfg',
                        default='experiments/cfgs/faster_rcnn_end2end.yml', type=str)
    parser.add_argument('--res', dest='res_file',
                        help='result file',
                        default='', type=str, required=True)
    args = parser.parse_args()
    return args


def run_test_net(gpu_id, caffemodel, prototxt, imdb_name, cfg_file):
    if cfg_file is not None:
        cfg_from_file(cfg_file)

    cfg.GPU_ID = gpu_id

    print('Using config:')
    pprint.pprint(cfg)

    while not os.path.exists(caffemodel):
        print('Waiting for {} to exist...'.format(caffemodel))
        time.sleep(10)

    caffe.set_mode_gpu()
    caffe.set_device(gpu_id)
    net = caffe.Net(prototxt, caffemodel, caffe.TEST)
    net.name = os.path.splitext(os.path.basename(caffemodel))[0]

    imdb = get_imdb(imdb_name)
    if not cfg.TEST.HAS_RPN:
        imdb.set_proposal_method(cfg.TEST.PROPOSAL_METHOD)

    n, _ = os.path.splitext(args.caffemodel)
    paths = splitall(n)

    proposal_prefix = paths[-1]

    return test_net(net, imdb, max_per_image=100, vis=False, proposal_prefix=proposal_prefix)


def run_test_nets(gpu_id, dir, model_files, prototxt, imdb_name, cfg_file, res_file):
    models = [line.rstrip('\n') for line in open(os.path.join(dir, model_files))]
    df_results = pd.DataFrame()
    for model in models:
        results = run_test_net(gpu_id, os.path.join(dir, model), prototxt, imdb_name, cfg_file)
        for result in results:
            result['file'] = model
        df_results = df_results.append(results, ignore_index=True)

    df_results.to_csv(os.path.join(dir, res_file))


if __name__ == '__main__':
    # args = parse_args()
    gpu_id = 0
    # dir = '/home/abhijitcbim/git/pedestrian-detector/output/faster_rcnn_end2end/train/backup'
    # model_files = 'test.txt'


    args = parse_args()

    print('Called with args:')
    print(args)

    run_test_nets(args.gpu_id, args.dir, args.model_files, args.prototxt, args.imdb_name, args.cfg_file, args.res_file)

    # run_test_net(gpu_id,caffemodel, prototxt, imdb_name, cfg_file)

[((43, 13, 43, 86), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((74, 4, 74, 22), 'pprint.pprint', 'pprint.pprint', ({(74, 18, 74, 21): 'cfg'}, {}), '(cfg)', False, 'import pprint\n'), ((80, 4, 80, 24), 'caffe.set_mode_gpu', 'caffe.set_mode_gpu', ({}, {}), '()', False, 'import caffe\n'), ((81, 4, 81, 28), 'caffe.set_device', 'caffe.set_device', ({(81, 21, 81, 27): 'gpu_id'}, {}), '(gpu_id)', False, 'import caffe\n'), ((82, 10, 82, 53), 'caffe.Net', 'caffe.Net', ({(82, 20, 82, 28): 'prototxt', (82, 30, 82, 40): 'caffemodel', (82, 42, 82, 52): 'caffe.TEST'}, {}), '(prototxt, caffemodel, caffe.TEST)', False, 'import caffe\n'), ((85, 11, 85, 30), 'datasets.factory.get_imdb', 'get_imdb', ({(85, 20, 85, 29): 'imdb_name'}, {}), '(imdb_name)', False, 'from datasets.factory import get_imdb\n'), ((89, 11, 89, 44), 'os.path.splitext', 'os.path.splitext', ({(89, 28, 89, 43): 'args.caffemodel'}, {}), '(args.caffemodel)', False, 'import time, os, sys\n'), ((94, 11, 94, 93), 'fast_rcnn.test.test_net', 'test_net', (), '', False, 'from fast_rcnn.test import test_net\n'), ((99, 17, 99, 31), 'pandas.DataFrame', 'pd.DataFrame', ({}, {}), '()', True, 'import pandas as pd\n'), ((26, 16, 26, 35), 'os.path.split', 'os.path.split', ({(26, 30, 26, 34): 'path'}, {}), '(path)', False, 'import time, os, sys\n'), ((69, 8, 69, 31), 'fast_rcnn.config.cfg_from_file', 'cfg_from_file', ({(69, 22, 69, 30): 'cfg_file'}, {}), '(cfg_file)', False, 'from fast_rcnn.config import cfg, cfg_from_file, cfg_from_list\n'), ((76, 14, 76, 40), 'os.path.exists', 'os.path.exists', ({(76, 29, 76, 39): 'caffemodel'}, {}), '(caffemodel)', False, 'import time, os, sys\n'), ((78, 8, 78, 22), 'time.sleep', 'time.sleep', ({(78, 19, 78, 21): '(10)'}, {}), '(10)', False, 'import time, os, sys\n'), ((106, 22, 106, 49), 'os.path.join', 'os.path.join', ({(106, 35, 106, 38): 'dir', (106, 40, 106, 48): 'res_file'}, {}), '(dir, res_file)', False, 'import time, os, sys\n'), ((83, 32, 83, 60), 'os.path.basename', 'os.path.basename', ({(83, 49, 83, 59): 'caffemodel'}, {}), '(caffemodel)', False, 'import time, os, sys\n'), ((101, 39, 101, 63), 'os.path.join', 'os.path.join', ({(101, 52, 101, 55): 'dir', (101, 57, 101, 62): 'model'}, {}), '(dir, model)', False, 'import time, os, sys\n'), ((98, 49, 98, 79), 'os.path.join', 'os.path.join', ({(98, 62, 98, 65): 'dir', (98, 67, 98, 78): 'model_files'}, {}), '(dir, model_files)', False, 'import time, os, sys\n')]