import numpy as np
import cv2
import os
import time
from tqdm import tqdm
from ultralytics import YOLO
from ultralytics.yolo.utils.plotting import Annotator

from yolov8_det import analysis_yolov8

model_yolo = YOLO("E:/Bank_files/Bank_02/model_files/all_labels.pt")


# 图像文件夹
def get_video_list(path):
    video_ext = [".mp4", ".avi",".MP4"]
    video_names = []
    for maindir, subdir, file_name_list in os.walk(path):
        for filename in file_name_list:
            apath = os.path.join(maindir, filename)
            ext = os.path.splitext(apath)[1]
            if ext in video_ext:
                video_names.append(apath)
    return video_names


# 截取裁剪需要的视频帧
def save_seg_video(video_name,frameToStart,frametoStop,videoWriter,bbox):

    cap = cv2.VideoCapture(video_name)
    count = 0

    while True:

        success, frame = cap.read()

        if success:

            count += 1
            if count <= frametoStop and count > frameToStart:  # 选取起始帧
                print('correct= ', count)
                
                #裁剪视频画面
                frame_target = frame[int(bbox[1]):int(bbox[3]), int(bbox[0]):int(bbox[2])]  # (split_height, split_width)

                videoWriter.write(frame_target)
                
        if not success or count >= frametoStop:
            break

        print('end')  



# 对视频的操作
def get_seg_video(video_file,video_save_path,dertTime):


    # 检查路径
    print("frame image save path：{}".format(video_save_path))
    os.makedirs(video_save_path, exist_ok=True)
    
    if os.path.isdir(video_file):
        files = get_video_list(video_file)
    else:
        files = [video_file]

    files.sort()
    video_num = len(files)

    for num in range(video_num):

        # 视频名字
        video_name = files[num]

        print(video_name)
        video_basename = os.path.basename(video_name).split('.')[0]

        cap = cv2.VideoCapture(video_name)
        #帧率
        fps = cap.get(cv2.CAP_PROP_FPS)

        success,frame = cap.read()

        count_fps = 0
        write_fps = 0
        dertTime = 2
        # 每段帧率
        dertF = dertTime * fps

        # while True:

        # 前后帧信息保存

        # result_list = []

        # count_result_num = 0

        while success:

                count_fps  += 1

                # 调用模型，逐帧检测
                results_img = analysis_yolov8(frame=frame,
                                            model_coco=model_yolo,
                                            confidence=0.1)
                

                # result_list.append({count_fps:results_img})

                # if len(result_list) == 5:

                #     result_list.clear()

                
                # num = len(results_img)

                # # 如果只检测到一个人
                if num == 1:

                    # 起始帧
                    write_fps = count_fps

                    stop_fps = write_fps + dertF

                    # 目标检测结果

                    bbox = list(results_img[0].values())[0]
                    w = bbox[2] -bbox[0]
                    h = bbox[3] -bbox[1]
                    size = [int(w),int(h)]
                    # 保存截取视频
                    video_name_save = video_save_path + '/' +  video_basename + '_' +str(write_fps) + '.avi'

                    videoWriter =cv2.VideoWriter(video_name_save,cv2.VideoWriter_fourcc('X','V','I','D'),fps,size)                    

                    save_seg_video(video_name,write_fps,stop_fps,videoWriter,bbox)

                    # result_dict = {count_fps:bbox}

                    # result_list.append(result_dict)
                    
                    # print(count_fps,write_fps,stop_fps,video_name,bbox)

                    break

                # if num == 0:

                #     continue

                # if num > 1:

                #     # print()
                #     pass



if __name__ == '__main__':  
     

    # 每个视频的时长（单位秒）
    dertTime = 5

    video = "E:/Bank_files/Bank_02/dataset/vlc_0711/0711-1.mp4"
    video_save = 'videos_codes_2'

    get_seg_video(video_file=video,video_save_path=video_save,dertTime=dertTime)