fcakyon commited on
Commit
7c89dbd
1 Parent(s): 445b8b8

Upload utils.py (#1)

Browse files

- Upload utils.py (6fb6cdbc323646518c48dc919d71c108516e66bc)

Files changed (1) hide show
  1. utils.py +36 -0
utils.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from pytube import YouTube
2
+ import numpy as np
3
+ from decord import VideoReader, cpu
4
+ import imageio
5
+
6
+
7
+ def download_youtube_video(url: str):
8
+ yt = YouTube(url)
9
+
10
+ streams = yt.streams.filter(file_extension='mp4')
11
+ file_path = streams[0].download()
12
+ return file_path
13
+
14
+
15
+ def sample_frame_indices(clip_len, frame_sample_rate, seg_len):
16
+ converted_len = int(clip_len * frame_sample_rate)
17
+ end_idx = np.random.randint(converted_len, seg_len)
18
+ start_idx = end_idx - converted_len
19
+ indices = np.linspace(start_idx, end_idx, num=clip_len)
20
+ indices = np.clip(indices, start_idx, end_idx - 1).astype(np.int64)
21
+ return indices
22
+
23
+
24
+ def sample_frames_from_video_file(file_path: str, num_frames: int = 16):
25
+ videoreader = VideoReader(file_path, num_threads=1, ctx=cpu(0))
26
+
27
+ # sample frames
28
+ videoreader.seek(0)
29
+ indices = sample_frame_indices(clip_len=num_frames, frame_sample_rate=4, seg_len=len(videoreader))
30
+ frames = videoreader.get_batch(indices).asnumpy()
31
+ return frames
32
+
33
+ def convert_frames_to_gif(frames):
34
+ converted_frames = frames.astype(np.uint8)
35
+ imageio.mimsave("frames.gif", converted_frames, fps=10)
36
+ return "frames.gif"