'''
author ：  Hatcher W huimin
Email： hcwang@semi.ac.cn
use: 目标检测视频处理
input:需要处理的视频路径
output：avi后缀的视频处理结果
'''
import glob
import os
import time

import torch
from PIL import Image
from vizer.draw import draw_boxes

from ssd.config import cfg
from ssd.data.datasets import COCODataset, VOCDataset
import argparse
import numpy as np

from ssd.data.transforms import build_transforms
from ssd.modeling.detector import build_detection_model
from ssd.utils import mkdir
from ssd.utils.checkpoint import CheckPointer
from detection import  main
from vediotoima import vedio_img
from imatovideo import images_to_video,images_to_video_big
from moviepy.editor import *
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '2'

def parse_args():

    parser = argparse.ArgumentParser(description='Detection Video to Video')
    parser.add_argument(
        "--config-file",
        default="ssd/config/hat-512.yaml",
        metavar="FILE",
        help="path to config file",
        type=str,
    )
    parser.add_argument("--ckpt", type=str, default="../data_MPKFNet/weigth/detection/detection.pth",
                        help="Trained weights.")
    parser.add_argument("--score_threshold", type=float, default=0.3)
    parser.add_argument("--video_name", default='000120210501080936.mp4', type=str, help='Video name.')
    parser.add_argument("--images_dir", default='../data_MPKFNet/result_raw/yang/', type=str,
                        help='Specify a image dir to do prediction.')
    parser.add_argument("--output_dir", default='../data_MPKFNet/result/yangd/', type=str,
                        help='Specify a image dir to save predicted images.')
    parser.add_argument("--dataset_type", default="voc", type=str,
                        help='Specify dataset type. Currently support voc and coco.')

    parser.add_argument(
        "opts",
        help="Modify config options using the command-line",
        default=None,
        nargs=argparse.REMAINDER,
    )
    args = parser.parse_args()
    return args


# if __name__ == '__main__':
#     args = parse_args()
#     path = "/data/Hatcher/output1/"
#     pathlists = os.listdir(path)
#     for pathlist in pathlists:
#         if os.path.exists('../data_MPKFNet/rawdata/' + pathlist + ".mp4"):
#             print(pathlist," is skipping")
#             continue
#         print(pathlist," is starting")
#         L = []
#         for root, dirs, files in os.walk(os.path.join(path,pathlist)):
#             files.sort()
#             for file in files:
#                 if os.path.splitext(file)[1] in [".avi",".mp4"]:
#                     filepath = os.path.join(root, file)
#                     video = VideoFileClip(filepath)
#                     L.append(video)
#         if len(L) == 0:
#             continue
#         final_clip = concatenate_videoclips(L)
#         final_clip.to_videofile(os.path.join('../data_MPKFNet/rawdata/', "{}.mp4".format(pathlist)), fps=25, remove_temp=False)
#         video_name_path = pathlist
#         video_name_path = video_name_path.replace(".MP4", "").replace(".mp4", "").replace(".MOV", "").replace(".mov", "")
#         RawVideoFilename = '../data_MPKFNet/rawdata/' + video_name_path + ".mp4"
#         RawVideoExtrationSavedpath = '../data_MPKFNet/result_raw/' + video_name_path
#         args.images_dir =  '../data_MPKFNet/result_raw/' + video_name_path
#         args.output_dir = "../data_MPKFNet/result/" + video_name_path + "Detection"
#         CombineVideoRawImgPath = "../data_MPKFNet/result/" + video_name_path + "Detection"
#         CombineVideoResultPath = '../data_MPKFNet/result_vedio11/'  + "{}.avi".format(video_name_path)
#
#         # vedio_img(RawVideoFilename, RawVideoExtrationSavedpath, 1)
#         # main(args)
#         # images_to_video(CombineVideoRawImgPath, CombineVideoResultPath)
#         if os.path.exists(RawVideoExtrationSavedpath):
#             print("skip video to image")
#             pass
#         else:
#             vedio_img(RawVideoFilename, RawVideoExtrationSavedpath, 1)
#         if os.path.exists(args.output_dir):
#             print("skip deal")
#             pass
#         else:
#             main(args)
#         images_to_video(CombineVideoRawImgPath, CombineVideoResultPath)


#zhendui danyi shipin
if __name__ == '__main__':
    args = parse_args()
    video_name_path = args.video_name
    video_name_path = video_name_path.replace(".MP4", "").replace(".mp4", "").replace(".MOV", "").replace(".mov", "")
    RawVideoFilename = '../data_MPKFNet/rawdata/1/'+args.video_name
    RawVideoExtrationSavedpath = '../data_MPKFNet/result_raw/' + video_name_path
    args.images_dir =  '../data_MPKFNet/result_raw/' + video_name_path
    args.output_dir = "../data_MPKFNet/result/" + video_name_path + "Detection"
    CombineVideoRawImgPath = "../data_MPKFNet/result/" + video_name_path + "Detection"
    CombineVideoResultPath = '../data_MPKFNet/result_vedio/' + video_name_path + "Detection.avi"
    if os.path.exists(RawVideoExtrationSavedpath):
        print("skip video to image")
        pass
    else:
        vedio_img(RawVideoFilename, RawVideoExtrationSavedpath, 1)
    if os.path.exists(args.output_dir):
        print("skip deal")
        pass
    else:
        main(args)
    images_to_video(CombineVideoRawImgPath, CombineVideoResultPath)