-
Notifications
You must be signed in to change notification settings - Fork 9
/
predict.py
81 lines (61 loc) · 2.93 KB
/
predict.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
import argparse
import imageio
import numpy as np
from tensorflow.keras.models import load_model
from PIL import Image, ImageOps
from tqdm import tqdm
from utils import draw_predictions, compute_metrics
def main(args):
video = imageio.get_reader(args.video)
n_frames = video.count_frames()
fps = video.get_meta_data()['fps']
frame_w, frame_h = video.get_meta_data()['size']
model = load_model(args.model, compile=False)
input_shape = model.input.shape[1:3]
# default RoI
if None in (args.rl, args.rt, args.rr, args.rb):
side = min(frame_w, frame_h)
args.rl = (frame_w - side) / 2
args.rt = (frame_h - side) / 2
args.rr = (frame_w + side) / 2
args.rb = (frame_h + side) / 2
crop = (args.rl, args.rt, args.rr, args.rb)
def preprocess(frame):
frame = Image.fromarray(frame)
eye = frame.crop(crop)
eye = ImageOps.grayscale(eye)
eye = eye.resize(input_shape)
return eye
def predict(eye):
eye = np.array(eye).astype(np.float32) / 255.0
eye = eye[None, :, :, None]
return model.predict(eye)
out_video = imageio.get_writer(args.output_video, fps=fps)
cropped = map(preprocess, video)
frames_and_predictions = map(lambda x: (x, predict(x)), cropped)
with open(args.output_csv, 'w') as out_csv:
print('frame,pupil-area,pupil-x,pupil-y,eye,blink', file=out_csv)
for idx, (frame, predictions) in enumerate(tqdm(frames_and_predictions, total=n_frames)):
pupil_map, tags = predictions
is_eye, is_blink = tags.squeeze()
(pupil_y, pupil_x), pupil_area = compute_metrics(pupil_map, thr=args.thr, nms=True)
row = [idx, pupil_area, pupil_x, pupil_y, is_eye, is_blink]
row = ','.join(list(map(str, row)))
print(row, file=out_csv)
img = draw_predictions(frame, predictions, thr=args.thr)
img = np.array(img)
out_video.append_data(img)
out_video.close()
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Predict on test video')
parser.add_argument('model', type=str, help='Path to model')
parser.add_argument('video', type=str, default='<video0>', help='Video file to process (use \'<video0>\' for webcam)')
parser.add_argument('-t', '--thr', type=float, default=0.5, help='Map Threshold')
parser.add_argument('-rl', type=int, help='RoI X coordinate of top left corner')
parser.add_argument('-rt', type=int, help='RoI Y coordinate of top left corner')
parser.add_argument('-rr', type=int, help='RoI X coordinate of right bottom corner')
parser.add_argument('-rb', type=int, help='RoI Y coordinate of right bottom corner')
parser.add_argument('-ov', '--output-video', default='predictions.mp4', help='Output video')
parser.add_argument('-oc', '--output-csv', default='pupillometry.csv', help='Output CSV')
args = parser.parse_args()
main(args)