import cv2
import glob
import os
from tqdm import tqdm

is_stati = True
time_split = [2, 3, 4, 5, 6, 12]

if __name__ == '__main__':
    if is_stati:
        path = os.path.join(os.getcwd(), 'result')
        path = path+'/*.mp4'
        pic_list = glob.glob(path)
        print("该目录下总共有：", len(pic_list), '个视频')
    else:
        if not os.path.exists('./result'):
            os.mkdir('./result')
        video_list = glob.glob(r'./*.mp4')
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')


        for video in tqdm(video_list):
            name = video.split('/')[-1]
            cap = cv2.VideoCapture(video)
            fps = int(cap.get(cv2.CAP_PROP_FPS))
            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

            name_list = name.split('-')
            assert len(name_list) == 4

            for time_stamp in time_split:
                if name_list[1] == '常规场景':
                    temp = name_list[1] + '_' + name_list[2] + '_' + str(time_stamp) + 's.mp4'
                else:
                    temp = '复杂场景_' + name_list[2] + '_' + str(time_stamp) + 's.mp4'
                video_path = os.path.join(os.getcwd(), 'result', temp)
                video_writer = cv2.VideoWriter(video_path, fourcc, fps, (width, height))

                start_frame = fps * 1
                end_frame = start_frame + fps * time_stamp

                cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)  # 设置开始帧
                current_frame = cap.get(cv2.CAP_PROP_POS_FRAMES)
                while current_frame < end_frame:  # 从start到end之间读取帧数
                    ret, frame = cap.read()  # 从开始帧开始读取，之后会从开始帧依次往后读取，直到退出循环
                    video_writer.write(frame)  # 利用'写入视频对象'写入帧
                    current_frame = cap.get(cv2.CAP_PROP_POS_FRAMES)  # 获取当前帧数
                video_writer.release()
            cap.release()