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")) | |