import os
import cv2

def create_dir(path):
    if not os.path.exists(path):
        os.makedirs(path)

def save_frames(video_path, output_dir, num_frames=16):
    # 创建输出文件夹
    create_dir(output_dir)
    
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    
    # 获取视频总帧数
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    frame_interval = max(1, total_frames // num_frames)
    
    current_frame = 0
    saved_frame_count = 0
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        
        # 每隔 frame_interval 保存一帧
        if current_frame % frame_interval == 0:
            frame_filename = os.path.join(output_dir, f'frame_{saved_frame_count:04d}.png')
            cv2.imwrite(frame_filename, frame)
            saved_frame_count += 1
            
            if saved_frame_count >= num_frames:
                break
        
        current_frame += 1
    
    cap.release()

def process_videos(root_folder):
    subfolders = ['sky', 'taichi', 'ucf101']
    
    for subfolder in subfolders:
        subfolder_path = os.path.join(root_folder, subfolder)
        
        for video_file in os.listdir(subfolder_path):
            if video_file.endswith('.mp4'):
                video_path = os.path.join(subfolder_path, video_file)
                video_name = os.path.splitext(video_file)[0]
                output_dir = os.path.join(subfolder_path, video_name)
                
                save_frames(video_path, output_dir, num_frames=16)

# 设置根文件夹路径
root_folder = '/data2/zhaoliangliang/sora_new/Open-Sora/sample_test'

# 处理视频
process_videos(root_folder)
