File size: 3,004 Bytes
eec42bd |
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 |
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)
|