Update utils.py

utils.py

@@ -1,36 +1,51 @@
+from pathlib import Path
 from pytube import YouTube
 import numpy as np
-from decord import VideoReader, cpu
+from decord import VideoReader
 import imageio
 
 
 def download_youtube_video(url: str):
     yt = YouTube(url)
 
-    streams = yt.streams.filter(file_extension='mp4')
+    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))
+def sample_frames_from_video_file(
+    file_path: str, num_frames: int = 16, frame_sampling_rate=1
+):
+    videoreader = VideoReader(file_path)
+    videoreader.seek(0)
 
     # sample frames
-    videoreader.seek(0)
-    indices = sample_frame_indices(clip_len=num_frames, frame_sample_rate=4, seg_len=len(videoreader))
+    start_idx = 0
+    end_idx = num_frames * frame_sampling_rate - 1
+    indices = np.linspace(start_idx, end_idx, num=num_frames, dtype=np.int64)
     frames = videoreader.get_batch(indices).asnumpy()
+
     return frames
 
-def convert_frames_to_gif(frames):
+
+def get_num_total_frames(file_path: str):
+    videoreader = VideoReader(file_path)
+    videoreader.seek(0)
+    return len(videoreader)
+
+
+def convert_frames_to_gif(frames, save_path: str = "frames.gif"):
     converted_frames = frames.astype(np.uint8)
-    imageio.mimsave("frames.gif", converted_frames, fps=10)
-    return "frames.gif"
+    Path(save_path).parent.mkdir(parents=True, exist_ok=True)
+    imageio.mimsave(save_path, converted_frames, fps=8)
+    return save_path
+
+
+def create_gif_from_video_file(
+    file_path: str,
+    num_frames: int = 16,
+    frame_sampling_rate: int = 1,
+    save_path: str = "frames.gif",
+):
+    frames = sample_frames_from_video_file(file_path, num_frames, frame_sampling_rate)
+    return convert_frames_to_gif(frames, save_path)