# Copyright (c) OpenMMLab. All rights reserved.
import argparse
from genericpath import exists
from re import T

import cv2
import torch
import mmcv
import datetime
import numpy as np
from imantics import Polygons, Mask
from shapely.geometry import Polygon
from mmdet.apis import inference_detector, init_detector
from demo_util import *
from Tracks import *
def parse_args():
    parser = argparse.ArgumentParser(description='MMDetection video demo')
    parser.add_argument('video', help='Video file')
    parser.add_argument('config', help='Config file')
    parser.add_argument('checkpoint', help='Checkpoint file')
    parser.add_argument(
        '--device', default='cuda:0', help='Device used for inference')
    parser.add_argument(
        '--score-thr', type=float, default=0.3, help='Bbox score threshold')
    parser.add_argument('--out', type=str, help='Output video file')
    parser.add_argument('--show', action='store_true', help='Show video')
    parser.add_argument(
        '--wait-time',
        type=float,
        default=1,
        help='The interval of show (s), 0 is block')
    args = parser.parse_args()
    return args  

def main():
    args = parse_args()
    assert args.out or args.show, \
        ('Please specify at least one operation (save/show the '
         'video) with the argument "--out" or "--show"')

    model = init_detector(args.config, args.checkpoint, device=args.device)

    video_reader = mmcv.VideoReader(args.video)
    video_writer = None
    if args.out:
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        video_writer = cv2.VideoWriter(
            args.out, fourcc, video_reader.fps,
            (video_reader.width, video_reader.height))

    for frame in mmcv.track_iter_progress(video_reader):
        result = inference_detector(model, frame)
        frame = model.show_result(frame, result, score_thr=args.score_thr)
        if args.show:
            cv2.namedWindow('video', 0)
            mmcv.imshow(frame, 'video', args.wait_time)
        if args.out:
            video_writer.write(frame)

    if video_writer:
        video_writer.release()
    cv2.destroyAllWindows()

def main_cv2():
    args = parse_args()
    assert args.out or args.show, \
        ('Please specify at least one operation (save/show the '
         'video) with the argument "--out" or "--show"')

    model = init_detector(args.config, args.checkpoint, device=args.device)
    print("Process:"+args.video)
    video_reader = cv2.VideoCapture(args.video)
    video_writer = None
    
    use_boundingbox = False
    use_mask = True

    #if "segms" in args.out:
    #    use_boundingbox = False
    #    use_mask = True
    save_dir = args.out
    save_dir = save_dir.replace( os.path.basename(save_dir),'')
    if not exists(save_dir):
        os.makedirs(save_dir)
    if args.out:
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        video_writer = cv2.VideoWriter(
            args.out, fourcc, video_reader.get(cv2.CAP_PROP_FPS),
            (int(video_reader.get(cv2.CAP_PROP_FRAME_WIDTH)),
            int(video_reader.get(cv2.CAP_PROP_FRAME_HEIGHT))))

    success,frame = video_reader.read()
    frame_index = 1
    tmer = tracks_manager()
    
    while success:
        #frame = cv2.flip(frame, 0)
        time0=datetime.datetime.now()          
        result = inference_detector(model, frame)
        time1=datetime.datetime.now()
        
        bbox_results, segm_results = result
        
        if use_boundingbox:
            tmer.update_with_bbox(bbox_results,frame_index)
            frame = show_result(frame, result, score_thr=args.score_thr,show_mask=False)
        elif use_mask:
            tmer.update_with_segm(segm_results,frame_index)
            frame = show_result(frame, result, score_thr=args.score_thr,show_bbox=False)        
        
        frame = tmer.vis(frame)
        #cv2.imwrite("/home/qilei/DATASETS/trans_drone/mmtest.jpg",frame)
        video_writer.write(frame)
        
        success,frame = video_reader.read()
        
        frame_index+=1
    if not use_boundingbox:
        if isinstance( args.out,str):
            tmer.save_results(args.out+".json")
        
        json2csv(args.out+".json")

if __name__ == '__main__':
    #main()
    main_cv2()
