import os
import threading

NUM_THREADS = 100
VIDEO_ROOT = '/data/whj/CIL_IN_VIDEO/PIVOT-main/datasets/path_data/UCF-101'  # 原始视频存放路径
FRAME_ROOT = '/data/whj/CIL_IN_VIDEO/PIVOT-main/datasets/path_frames/UCF-101'  # 提取帧存放路径


def split(l, n):
    """将列表 l 分割为大小为 n 的子列表"""
    for i in range(0, len(l), n):
        yield l[i:i + n]


def extract(video_path, frame_dir, tmpl='%06d.jpg'):
    """调用 ffmpeg 提取视频帧"""
    cmd = 'ffmpeg -i \"{}\" -threads 1 -vf scale=-1:256 -q:v 0 \"{}/{}\"'.format(
        video_path, frame_dir, tmpl
    )
    os.system(cmd)


def target(video_list):
    """处理视频列表中的每个视频"""
    for video_path, class_dir in video_list:
        # 获取视频文件名（不含后缀）
        video_name = os.path.splitext(os.path.basename(video_path))[0]
        # 提取帧存放的路径
        frame_dir = os.path.join(FRAME_ROOT, class_dir, video_name)
        os.makedirs(frame_dir, exist_ok=True)  # 确保目录存在

        # 提取视频帧
        extract(video_path, frame_dir)


def get_video_list(video_root):
    """递归获取所有视频文件路径及对应的类名"""
    video_list = []
    for root, dirs, files in os.walk(video_root):
        # 获取当前目录的类名
        class_dir = os.path.relpath(root, video_root)
        for file in files:
            # 判断是否是视频文件（根据扩展名过滤）
            if file.endswith(('.avi', '.mp4', '.mkv', '.webm')):
                video_path = os.path.join(root, file)
                video_list.append((video_path, class_dir))  # 保存视频路径及类名
    return video_list


if __name__ == '__main__':
    # 检查视频根路径
    if not os.path.exists(VIDEO_ROOT):
        raise ValueError('请确认视频路径 VIDEO_ROOT 是否正确')
    # 确保帧存放路径存在
    if not os.path.exists(FRAME_ROOT):
        os.makedirs(FRAME_ROOT)

    # 获取所有视频文件列表
    video_list = get_video_list(VIDEO_ROOT)
    # 将视频列表分割为多个子列表
    splits = list(split(video_list, NUM_THREADS))

    # 多线程处理视频提取
    threads = []
    for i, split in enumerate(splits):
        thread = threading.Thread(target=target, args=(split,))
        thread.start()
        threads.append(thread)

    # 等待所有线程完成
    for thread in threads:
        thread.join()

    print("视频帧提取完成！")