import os
from moviepy.editor import VideoFileClip
from tqdm import tqdm

def extract_frames(video_path, output_dir, frames_per_set=40):
    # Load the video
    clip = VideoFileClip(video_path)
    
    # Create the output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)

    frame_count = 0
    set_count = 0
    current_set = None
    read_count = 0

    for frame in tqdm(clip.iter_frames(fps=clip.fps, dtype="uint8"), desc="Extracting frames"):

        read_count+=1
        if read_count%2==0:continue # skip every other frame
        
        if frame_count % frames_per_set == 0:
            set_count += 1
            current_set = os.path.join(output_dir, f"set_{set_count}")
            os.makedirs(current_set, exist_ok=True)

        frame_path = os.path.join(current_set, f"frame_{frame_count + 1:08d}.jpg")
        VideoFileClip.save_frame(clip, frame_path, t=frame_count / clip.fps)
        
        frame_count += 1

    clip.close()
    print(f"Extracted {frame_count} frames into {set_count} sets in {output_dir}")

# Example usage:
video_file = "nature.webm"  # Works for webm, mp4, etc.
output_directory = "nature_frames/"

extract_frames(video_file, output_directory)