-
Notifications
You must be signed in to change notification settings - Fork 0
/
vid2img_emoreact_mask.py
executable file
·55 lines (42 loc) · 1.37 KB
/
vid2img_emoreact_mask.py
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
import os
import threading
import glob
import tqdm
import sys
NUM_THREADS = 10
def split(l, n):
"""Yield successive n-sized chunks from l."""
for i in range(0, len(l), n):
yield l[i:i + n]
#-vf scale=-1:256
def extract(video, frame_folder, tmpl='%06d.jpg'):
cmd = 'ffmpeg -i \"{}\" -threads 1 -q:v 0 \"{}/%06d.jpg\" -loglevel quiet'.format(video, frame_folder)
os.system(cmd)
def target(video_list):
for video in video_list:
frame_folder = video[:-4]+'_frames_full/'
if not os.path.exists(frame_folder):
os.makedirs(frame_folder)
extract(video, frame_folder)
def main(VIDEO_ROOT):
if not os.path.exists(VIDEO_ROOT):
raise ValueError('No directory: ' + VIDEO_ROOT)
video_list = glob.glob(VIDEO_ROOT + '*.mp4')
splits = list(split(video_list, NUM_THREADS))
threads = []
for i, spl in enumerate(splits):
thread = threading.Thread(target=target, args=(spl,))
thread.start()
threads.append(thread)
for thread in threads:
thread.join()
if __name__ == '__main__':
train_root = "/gpu-data2/nkeg/EmoReact/Data_mask/Train_mask/"
val_root = "/gpu-data2/nkeg/EmoReact/Data_mask/Validation_mask/"
test_root = "/gpu-data2/nkeg/EmoReact/Data_mask/Test_mask/"
print("Extracting training set video frames...")
main(train_root)
print("Extracting validation set video frames...")
main(val_root)
print("Extracting test set video frames...")
main(test_root)