-
Notifications
You must be signed in to change notification settings - Fork 134
/
load_videos.py
101 lines (88 loc) · 4.31 KB
/
load_videos.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
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
import numpy as np
import pandas as pd
import imageio
import os
import subprocess
from multiprocessing import Pool
from itertools import cycle
import warnings
import glob
import time
from tqdm import tqdm
from util import save
from argparse import ArgumentParser
from skimage import img_as_ubyte
from skimage.transform import resize
warnings.filterwarnings("ignore")
DEVNULL = open(os.devnull, 'wb')
def download(video_id, args):
video_path = os.path.join(args.video_folder, video_id + ".mp4")
subprocess.call([args.youtube, '-f', "''best/mp4''", '--write-auto-sub', '--write-sub',
'--sub-lang', 'en', '--skip-unavailable-fragments',
"https://www.youtube.com/watch?v=" + video_id, "--output",
video_path], stdout=DEVNULL, stderr=DEVNULL)
return video_path
def run(data):
video_id, args = data
if not os.path.exists(os.path.join(args.video_folder, video_id.split('#')[0] + '.mp4')):
download(video_id.split('#')[0], args)
if not os.path.exists(os.path.join(args.video_folder, video_id.split('#')[0] + '.mp4')):
print ('Can not load video %s, broken link' % video_id.split('#')[0])
return
reader = imageio.get_reader(os.path.join(args.video_folder, video_id.split('#')[0] + '.mp4'))
fps = reader.get_meta_data()['fps']
df = pd.read_csv(args.metadata)
df = df[df['video_id'] == video_id]
all_chunks_dict = [{'start': df['start'].iloc[j], 'end': df['end'].iloc[j],
'bbox': list(map(int, df['bbox'].iloc[j].split('-'))), 'frames':[]} for j in range(df.shape[0])]
ref_fps = df['fps'].iloc[0]
ref_height = df['height'].iloc[0]
ref_width = df['width'].iloc[0]
partition = df['partition'].iloc[0]
try:
for i, frame in enumerate(reader):
for entry in all_chunks_dict:
if (i * ref_fps >= entry['start'] * fps) and (i * ref_fps < entry['end'] * fps):
left, top, right, bot = entry['bbox']
left = int(left / (ref_width / frame.shape[1]))
top = int(top / (ref_height / frame.shape[0]))
right = int(right / (ref_width / frame.shape[1]))
bot = int(bot / (ref_height / frame.shape[0]))
crop = frame[top:bot, left:right]
if args.image_shape is not None:
crop = img_as_ubyte(resize(crop, args.image_shape, anti_aliasing=True))
entry['frames'].append(crop)
except imageio.core.format.CannotReadFrameError:
None
for entry in all_chunks_dict:
if 'person_id' in df:
first_part = df['person_id'].iloc[0] + "#"
else:
first_part = ""
first_part = first_part + '#'.join(video_id.split('#')[::-1])
path = first_part + '#' + str(entry['start']).zfill(6) + '#' + str(entry['end']).zfill(6) + '.mp4'
save(os.path.join(args.out_folder, partition, path), entry['frames'], args.format)
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument("--video_folder", default='youtube-taichi', help='Path to youtube videos')
parser.add_argument("--metadata", default='taichi-metadata-new.csv', help='Path to metadata')
parser.add_argument("--out_folder", default='taichi-png', help='Path to output')
parser.add_argument("--format", default='.png', help='Storing format')
parser.add_argument("--workers", default=1, type=int, help='Number of workers')
parser.add_argument("--youtube", default='./youtube-dl', help='Path to youtube-dl')
parser.add_argument("--image_shape", default=(256, 256), type=lambda x: tuple(map(int, x.split(','))),
help="Image shape, None for no resize")
args = parser.parse_args()
if not os.path.exists(args.video_folder):
os.makedirs(args.video_folder)
if not os.path.exists(args.out_folder):
os.makedirs(args.out_folder)
for partition in ['test', 'train']:
if not os.path.exists(os.path.join(args.out_folder, partition)):
os.makedirs(os.path.join(args.out_folder, partition))
df = pd.read_csv(args.metadata)
video_ids = set(df['video_id'])
pool = Pool(processes=args.workers)
args_list = cycle([args])
for chunks_data in tqdm(pool.imap_unordered(run, zip(video_ids, args_list))):
None