-
Notifications
You must be signed in to change notification settings - Fork 134
/
crop_nemo.py
86 lines (70 loc) · 3.42 KB
/
crop_nemo.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
import numpy as np
import pandas as pd
import imageio
from skimage.transform import resize
from argparse import ArgumentParser
from skimage import img_as_ubyte
import os
import subprocess
from multiprocessing import Process
import warnings
import glob
from tqdm import tqdm
import face_alignment
from util import bb_intersection_over_union, join, scheduler, crop_bbox_from_frames, save
warnings.filterwarnings("ignore")
TEST_PERSONS = {'133', '492', '174', '105', '445', '166', '525', '162', '447', '336', '071', '414', '116',
'148', '502', '225', '205', '093', '141', '004', '456', '263', '418', '483', '265', '450', '201',
'304', '505', '536', '510', '172', '112', '400', '270', '215', '155', '553', '343', '176', '213'}
REF_FPS = 50
def extract_bbox(frame, fa):
bbox = fa.face_detector.detect_from_image(frame[..., ::-1])[0]
return bbox
def store(frame_list, tube_bbox, video_id, args):
out, final_bbox = crop_bbox_from_frames(frame_list, tube_bbox, min_frames=0,
image_shape=args.image_shape, min_size=0,
increase_area=args.increase)
if out is None:
return []
name = video_id
person_id = video_id.split('_')[0]
partition = 'test' if person_id in TEST_PERSONS else 'train'
save(os.path.join(args.out_folder, partition, name), out, args.format)
return [{'bbox': '-'.join(map(str, final_bbox)), 'start': 0, 'end': len(frame_list), 'fps': REF_FPS,
'video_id': video_id, 'height': frame_list[0].shape[0], 'width': frame_list[0].shape[1], 'partition': partition}]
def process_video(video_id, args):
fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, flip_input=False)
video_path = os.path.join(args.in_folder, video_id)
reader = imageio.get_reader(video_path)
tube_bbox = None
frame_list = []
for i, frame in enumerate(reader):
if i == 0:
bbox = extract_bbox(resize(frame, (360, 640), preserve_range=True), fa)
bbox = bbox * 3
#left, top, right, bot, _ = bbox
tube_bbox = bbox[:-1]
frame_list.append(frame)
return store(frame_list, tube_bbox, video_id, args)
def run(params):
video_id, device_id, args = params
os.environ['CUDA_VISIBLE_DEVICES'] = device_id
return process_video(video_id, args)
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument("--in_folder", default = 'nemo-videos')
parser.add_argument("--out_folder", default = 'nemo-256')
parser.add_argument("--increase", default = 0.1, type=float, help='Increase bbox by this amount')
parser.add_argument("--format", default='.png', help='Store format (.png, .mp4)')
parser.add_argument("--chunks_metadata", default='nemo-metadata.csv', help='Path to store metadata')
parser.add_argument("--image_shape", default=(256, 256), type=lambda x: tuple(map(int, x.split(','))),
help="Image shape")
parser.add_argument("--workers", default=1, type=int, help='Number of parallel workers')
parser.add_argument("--device_ids", default="0", help="Names of the devices comma separated.")
args = parser.parse_args()
if not os.path.exists(args.out_folder):
os.makedirs(args.out_folder)
os.makedirs(args.out_folder + '/train')
os.makedirs(args.out_folder + '/test')
ids = sorted(os.listdir(args.in_folder))
scheduler(ids, run, args)