import project_path
import argparse
from infer_frames import main as infer
import sys
sys.path.append('..')
sys.path.append('../caltech-fish-counting')

from evaluate import evaluate

class Object(object):
    pass

def main(args):

    infer_args = Object()
    infer_args.metadata = "../frames/metadata"
    infer_args.frames = "../frames/images"
    infer_args.output = "../frames/result"
    infer_args.weights = "models/v5m_896_300best.pt"

    config = {
        'conf_threshold': float(args.conf_threshold),
        'nms_iou': float(args.nms_iou),
        'min_length': float(args.min_length),
        'max_age': int(args.max_age),
        'iou_threshold': float(args.iou_threshold),
        'min_hits': int(args.min_hits)
    }

    infer(infer_args, config=config, verbose=False)

    evaluate("../frames/result_testing", "../frames/MOT", "../frames/metadata", "tracker", False)


def argument_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument("--conf_threshold", default=0.3, help="Config object. Required.")
    parser.add_argument("--nms_iou", default=0.3, help="Config object. Required.")
    parser.add_argument("--min_length", default=0.3, help="Config object. Required.")
    parser.add_argument("--max_age", default=20, help="Config object. Required.")
    parser.add_argument("--iou_threshold", default=0.01, help="Config object. Required.")
    parser.add_argument("--min_hits", default=11, help="Config object. Required.")
    return parser

if __name__ == "__main__":
    args = argument_parser().parse_args()
    main(args)