import cv2
import os
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def extract_frames(url_path, output_dir) -> int :
    '''
    Acts as initial feed into the SuperSlomo Model
    The Frames are stored in an output directory which is then loaded into the SuperSlomo Model.
    :param url_path:
    :param output_dir:
    :return: None
    '''
    os.makedirs(output_dir, exist_ok=True)
    frame_count = 0
    cap = cv2.VideoCapture(url_path)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    while cap.isOpened():
        ret, frame = cap.read()  # frame is a numpy array
        if not ret:
            break
        frame_name = f"frame_{frame_count}.png"
        frame_count += 1
        cv2.imwrite(os.path.join(output_dir, frame_name), frame)
    cap.release()
    return fps


def downsample(video_path, output_dir, target_fps):
    pass



if __name__ == "__main__":  # sets the __name__ variable to __main__ for this script

    print(extract_frames("Test.mp4", "output"))