File size: 1,786 Bytes
08f6d0e 642ebc0 816f401 08f6d0e 642ebc0 816f401 642ebc0 08f6d0e 642ebc0 816f401 08f6d0e 642ebc0 816f401 642ebc0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
import cv2
import os
from PIL import Image
from torchvision.transforms import transforms, ToTensor
from torch import tensor
from torchvision.transforms import ToPILImage,Resize
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_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
def load_frames(path,size=(128,128)) -> tensor: # converts PIL image to tensor on the GPU
image = Image.open(path).convert('RGB')
tensor = ToTensor()
resized_image=Resize(size)(image)
return tensor(resized_image).unsqueeze(0).to(device)
def save_frames(Tensor,output_path)->None: # Tensor to image
'''
Used to Save the Interpolated frame into the output directory.
:param Tensor:
:param output_path:
:return:
'''
transform=ToPILImage()
image=Tensor.squeeze(0).cpu()
image=transform(image)
image.save(output_path)
if __name__ == "__main__": # sets the __name__ variable to __main__ for this script
extract_frames("Test.mp4", "output")
|