import cv2
import os
from tqdm import tqdm
from argparse import ArgumentParser, Namespace

parser = ArgumentParser()
parser.add_argument("--video_folder", default='demo_videos', type=str, help="Folder path of your videos")
parser.add_argument("--output_frame_folder", default='output_frames', type=str, help="Folder path of extracted frames")
args = parser.parse_args()

all_video = args.video_folder  # folder path of your videos
video_list = [os.path.join(all_video, path) for path in os.listdir(all_video)]

output_dir = args.output_frame_folder # output folder path of extracted frames
os.makedirs(output_dir, exist_ok=True)

print('Extracting frames from videos...')
for video_path in tqdm(video_list):
    fname = video_path.split('/')[-1][:-4]
    cap = cv2.VideoCapture(video_path)

    if not cap.isOpened():
        print("Error: Could not open video.")
    
    else:
        fps = cap.get(cv2.CAP_PROP_FPS)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        
        num_frames = 15 # default number of frames extracted in each video
        interval = total_frames // num_frames

        sampled_frames = [i * interval for i in range(num_frames)]

        for saved_frame_count, frame_idx in enumerate(sampled_frames):
            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
            ret, frame = cap.read()
            if ret:               
                frame_filename = os.path.join(output_dir,'{}_{}.jpg'.format(fname, str(saved_frame_count).zfill(3)))
                cv2.imwrite(frame_filename, frame)
                
            else:
                print(f"Frame {frame_idx} could not be read")

        cap.release()