|
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] |
|
|
|
|
|
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) |
|
|
