# 將四個(gè)文件夾中的多個(gè)視頻文件轉(zhuǎn)換為圖片
import cv2
import os


number_in_dir = []  # 可以去掉
def video2image(video_path, image_save_dir):
    # 1. 將圖片的存儲路徑加上源文件的文件夾名，如'Coffee_room_01'等。
    if os.path.exists(video_path):  # 判斷源路徑是否正確
        print(video_path + '\t ok')
        if video_path.split('\\')[-2] != 'Videos':  # 這個(gè)根據(jù)我文件路徑自行添加的
            image_save_dir += video_path.split('\\')[-2] + '\\'
        else:
            image_save_dir += video_path.split('\\')[-3] + '\\'

        number_in_dir.append(len(os.listdir(video_path)))  # 可以去掉
    else:
        print(video_path + ' \033[0;37;41merror\033[0m')
        return

    # 2. 依次讀取源文件里的文件，如果后綴名是‘a(chǎn)vi'或 'MP3'，則創(chuàng)建一個(gè)關(guān)于文件名的子文件夾
    files_in_video_path_list = os.listdir(video_path)
    print(len(files_in_video_path_list))  # 獲取當(dāng)前文件夾下文件個(gè)數(shù)，用來判斷是否有其他類型的文件
    file_count = 0  # 用于統(tǒng)計(jì)個(gè)數(shù)，驗(yàn)證是否全為視頻文件，會與len(files_in_video_path_list)進(jìn)行比較

    for file_in_video_path in files_in_video_path_list:
        file_name = os.path.basename(file_in_video_path)
        if file_name.split('.')[-1] == 'avi' or file_name.split('.')[-1] == 'mp3':
            file_count += 1  # 視頻文件數(shù)+1
            every_video_save_dir = file_name.split('.')[0] + '\\'
            if not os.path.exists(image_save_dir + every_video_save_dir):  # 創(chuàng)建屬于相應(yīng)文件夾的存儲路徑
                os.makedirs(image_save_dir + every_video_save_dir)
        else:
            print(' \033[0;37;41merror\033[0m')
            break

        # 3. 打印正在處理文件的序號和他的文件名，并開始轉(zhuǎn)換
        print(str(file_count) + ':' + file_name)
        cap = cv2.VideoCapture(video_path + file_name)

        flag = cap.isOpened()
        if not flag:
            print("open" + video_path + file_name + "error!")

        frame_count = 0  # 給每一幀標(biāo)號
        while True:
            frame_count += 1
            flag, frame = cap.read()
            if not flag:  # 如果已經(jīng)讀取到最后一幀則退出
                break
            if os.path.exists(image_save_dir + every_video_save_dir + str(frame_count) + '.jpg'):  # 在源視頻不變的情況下，如果已經(jīng)創(chuàng)建，則跳過
                break
            cv2.imwrite(image_save_dir + every_video_save_dir + str(frame_count) + '.jpg', frame)

        cap.release()
        print(file_name + ' save to ' + image_save_dir + every_video_save_dir + 'finished ')  # 表示一個(gè)視頻片段已經(jīng)轉(zhuǎn)換完成
    if file_count != len(files_in_video_path_list):
        print(' \033[0;37;41Maybe loss something\033[0m')


if __name__ == '__main__':
    # 需要轉(zhuǎn)換的文件路徑列表(自定義修改）
    video_path_list = ['E:\\Datasets_FallDetection\\2 LE2I\Coffee_room_01\\Coffee_room_01\\Videos\\',
                       'E:\\Datasets_FallDetection\\2 LE2I\Home_01\\Home_01\Videos\\',
                       'E:\\Datasets_FallDetection\\2 LE2I\Lecture_room\\Lecture room\\',
                       'E:\\Datasets_FallDetection\\2 LE2I\\Office\\Office\\']

    # 預(yù)期存儲在的主文件夾，即'result'文件夾
    image_save_dir = '.\\result\\'

    # 進(jìn)行轉(zhuǎn)換
    for video_path in video_path_list:
        video2image(video_path, image_save_dir)