Upload utils.py

utils.py

@@ -0,0 +1,36 @@
+from pytube import YouTube
+import numpy as np
+from decord import VideoReader, cpu
+import imageio
+
+
+def download_youtube_video(url: str):
+    yt = YouTube(url)
+
+    streams = yt.streams.filter(file_extension='mp4')
+    file_path = streams[0].download()
+    return file_path
+
+
+def sample_frame_indices(clip_len, frame_sample_rate, seg_len):
+    converted_len = int(clip_len * frame_sample_rate)
+    end_idx = np.random.randint(converted_len, seg_len)
+    start_idx = end_idx - converted_len
+    indices = np.linspace(start_idx, end_idx, num=clip_len)
+    indices = np.clip(indices, start_idx, end_idx - 1).astype(np.int64)
+    return indices
+
+
+def sample_frames_from_video_file(file_path: str, num_frames: int = 16):
+    videoreader = VideoReader(file_path, num_threads=1, ctx=cpu(0))
+
+    # sample frames
+    videoreader.seek(0)
+    indices = sample_frame_indices(clip_len=num_frames, frame_sample_rate=4, seg_len=len(videoreader))
+    frames = videoreader.get_batch(indices).asnumpy()
+    return frames
+
+def convert_frames_to_gif(frames):
+    converted_frames = frames.astype(np.uint8)
+    imageio.mimsave("frames.gif", converted_frames, fps=10)
+    return "frames.gif"