fcakyon/zero-shot-video-classification
52
1from pathlib import Path2from pytube import YouTube3import numpy as np4from decord import VideoReader5import imageio6 7 8def download_youtube_video(url: str):9 yt = YouTube(url)10 11 streams = yt.streams.filter(file_extension="mp4")12 file_path = streams[0].download()13 return file_path14 15 16def sample_frames_from_video_file(17 file_path: str, num_frames: int = 16, frame_sampling_rate=118):19 videoreader = VideoReader(file_path)20 videoreader.seek(0)21 22 # sample frames23 start_idx = 024 end_idx = num_frames * frame_sampling_rate - 125 indices = np.linspace(start_idx, end_idx, num=num_frames, dtype=np.int64)26 frames = videoreader.get_batch(indices).asnumpy()27 28 return frames29 30 31def get_num_total_frames(file_path: str):32 videoreader = VideoReader(file_path)33 videoreader.seek(0)34 return len(videoreader)35 36 37def convert_frames_to_gif(frames, save_path: str = "frames.gif"):38 converted_frames = frames.astype(np.uint8)39 Path(save_path).parent.mkdir(parents=True, exist_ok=True)40 imageio.mimsave(save_path, converted_frames, fps=8)41 return save_path42 43 44def create_gif_from_video_file(45 file_path: str,46 num_frames: int = 16,47 frame_sampling_rate: int = 1,48 save_path: str = "frames.gif",49):50 frames = sample_frames_from_video_file(file_path, num_frames, frame_sampling_rate)51 return convert_frames_to_gif(frames, save_path)52 