#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2019/10/30 11:55
# @Author  : Scheaven
# @File    :  demo.py
# @description:
import argparse
from lib.core.ds_tracker import human_tracker
from lib.core.yolo import YOLO
from lib.core import generate_detections as gdet
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

######################paraters######################
def parse_args():
    parser = argparse.ArgumentParser(description="Deep SORT")
    parser.add_argument("-i", "--in_type", help="camera or video",
        default='video', required=False)
    parser.add_argument("--c_in", help="camera ",
        default='rtsp://admin:a1234567@192.168.5.32:554/h264/ch1/main/av_stream', required=False)
    parser.add_argument("--v_in", help="video",
        default='../cs01.avi', required=False)
    return parser.parse_args()

if __name__ == '__main__':
    args = parse_args()
    model_filename = 'model_dump/mars-small128.pb'
    encoder = gdet.create_box_encoder(model_filename, batch_size=1)
    human_tracker(YOLO(), encoder, args)