headscratchertm's picture
works fine this shit is so ass
9b56f03
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"))