From 4782ee8ab73189a26aae3419d89e761705526a1d Mon Sep 17 00:00:00 2001 From: Pablo Prietz Date: Wed, 14 Mar 2018 14:47:46 +0100 Subject: [PATCH] Add iMotions Exporter --- pupil_src/shared_modules/imotions_exporter.py | 264 ++++++++++++++++++ 1 file changed, 264 insertions(+) create mode 100644 pupil_src/shared_modules/imotions_exporter.py diff --git a/pupil_src/shared_modules/imotions_exporter.py b/pupil_src/shared_modules/imotions_exporter.py new file mode 100644 index 0000000000..94ec355768 --- /dev/null +++ b/pupil_src/shared_modules/imotions_exporter.py @@ -0,0 +1,264 @@ +''' +(*)~--------------------------------------------------------------------------- +Pupil - eye tracking platform +Copyright (C) 2012-2017 Pupil Labs + +Distributed under the terms of the GNU +Lesser General Public License (LGPL v3.0). +See COPYING and COPYING.LESSER for license details. +---------------------------------------------------------------------------~(*) +''' + +import os +import csv +import av +import multiprocessing as mp +from glob import glob +from shutil import copy2 +from fractions import Fraction + +import numpy as np +from pyglui import ui +from plugin import Analysis_Plugin_Base +from video_capture import File_Source, EndofVideoFileError +from methods import denormalize +import background_helper as bh +import csv_utils + +import logging +logger = logging.getLogger(__name__) + + +__version__ = 2 + + +class Empty(object): + pass + + +def bitrate(video_height, fps): + ''' + Returns bitrate given these recommendations: + https://support.google.com/youtube/answer/1722171?hl=en + ''' + if video_height < 420: + bitrate = 1e6 + elif video_height < 600: + bitrate = 2.5e6 + elif video_height < 900: + bitrate = 5e6 + else: + bitrate = 8e6 + + if fps > 40: + bitrate *= 1.5 + return bitrate * 50 + + +def export_undistorted_h264(distorted_video_loc, target_video_loc, export_range): + yield "Converting scene video", .1 + capture = File_Source(Empty(), distorted_video_loc) + if not capture.initialised: + yield "Converting scene video failed", 0. + return + + update_rate = 10 + start_time = None + time_base = Fraction(1, 65535) + average_fps = int(len(capture.timestamps) / (capture.timestamps[-1] - capture.timestamps[0])) + + target_container = av.open(target_video_loc, 'w') + video_stream = target_container.add_stream('mpeg4', 1/time_base) + video_stream.bit_rate = bitrate(capture.frame_size[1], average_fps) + video_stream.bit_rate_tolerance = video_stream.bit_rate / 5 + video_stream.thread_count = max(1, mp.cpu_count() - 1) + video_stream.width, video_stream.height = capture.frame_size + + av_frame = av.VideoFrame(*capture.frame_size, 'bgr24') + av_frame.time_base = time_base + + capture.seek_to_frame(export_range[0]) + next_update_idx = export_range[0] + update_rate + while True: + try: + frame = capture.get_frame() + except EndofVideoFileError: + break + + if frame.index > export_range[1]: + break + + if start_time is None: + start_time = frame.timestamp + + undistorted_img = capture.intrinsics.undistort(frame.img) + av_frame.planes[0].update(undistorted_img) + av_frame.pts = int((frame.timestamp - start_time) / time_base) + + packet = video_stream.encode(av_frame) + if packet: + target_container.mux(packet) + + if capture.current_frame_idx >= next_update_idx: + progress = ((capture.current_frame_idx - export_range[0]) / + (export_range[1] - export_range[0])) * .9 + .1 + yield "Converting scene video", progress * 100. + next_update_idx += update_rate + + while True: # flush encoder + packet = video_stream.encode() + if packet: + target_container.mux(packet) + else: + break + + target_container.close() + capture.cleanup() + yield "Converting scene video completed", 1. * 100. + + +class iMotions_Exporter(Analysis_Plugin_Base): + '''iMotions Gaze and Video Exporter + + All files exported by this plugin are saved to a subdirectory within + the export directory called "iMotions". The gaze data will be written + into a file called "gaze.tlv" and the undistored scene video will be + saved in a file called "scene.mp4". + + The gaze.tlv file is a tab-separated CSV file with the following fields: + GazeTimeStamp: Timestamp of the gaze point, unit: seconds + MediaTimeStamp: Timestamp of the scene frame to which the gaze point + was correlated to, unit: seconds + MediaFrameIndex: Index of the scene frame to which the gaze point was + correlated to + Gaze3dX: X position of the 3d gaze point (the point the subject looks + at) in the scene camera coordinate system + Gaze3dY: Y position of the 3d gaze point + Gaze3dZ: Z position of the 3d gaze point + Gaze2dX: undistorted gaze pixel postion, X coordinate, unit: pixels + Gaze2dX: undistorted gaze pixel postion, Y coordinate, unit: pixels + PupilDiaLeft: Left pupil diameter, 0.0 if not available, unit: millimeters + PupilDiaRight: Right pupil diameter, 0.0 if not available, unit: millimeters + Confidence: Value between 0 and 1 indicating the quality of the gaze + datum. It depends on the confidence of the pupil detection + and the confidence of the 3d model. Higher values are good. + ''' + icon_chr = 'iM' + + def __init__(self, g_pool): + super().__init__(g_pool) + self.export_task = None + self.status = 'Not exporting' + self.progress = 0. + self.output = 'Not set yet' + logger.info('iMotions Exporter has been launched.') + + def on_notify(self, notification): + if notification['subject'] == "should_export": + self.export_data(notification['range'], notification['export_dir']) + + def init_ui(self): + self.add_menu() + self.menu.label = 'iMotions Exporter' + self.menu.append(ui.Text_Input('status', self, label='Status', setter=lambda _: None)) + self.menu.append(ui.Text_Input('output', self, label='Last export', setter=lambda _: None)) + self.menu.append(ui.Slider('progress', self, label='Progress')) + self.menu[-1].read_only = True + self.menu[-1].display_format = '%.0f%%' + self.menu.append(ui.Button('Cancel export', self.cancel)) + + def cancel(self): + if self.export_task: + self.export_task.cancel() + + def deinit_ui(self): + self.remove_menu() + + def cleanup(self): + self.cancel() + + def export_data(self, export_range, export_dir): + rec_start = self.get_recording_start_date() + im_dir = os.path.join(export_dir, 'iMotions_{}'.format(rec_start)) + os.makedirs(im_dir, exist_ok=True) + user_warned_3d_only = False + self.output = im_dir + logger.info('Exporting to {}'.format(im_dir)) + + if self.export_task: + self.export_task.cancel() + + distorted_video_loc = [f for f in glob(os.path.join(self.g_pool.rec_dir, "world.*")) + if os.path.splitext(f)[-1] in ('.mp4', '.mkv', '.avi', '.mjpeg')][0] + target_video_loc = os.path.join(im_dir, 'scene.mp4') + generator_args = (distorted_video_loc, target_video_loc, export_range) + self.export_task = bh.Task_Proxy('iMotions Video Export', export_undistorted_h264, + args=generator_args) + + info_src = os.path.join(self.g_pool.rec_dir, 'info.csv') + info_dest = os.path.join(im_dir, 'iMotions_info.csv') + copy2(info_src, info_dest) # copy info.csv file + + with open(os.path.join(im_dir, 'gaze.tlv'), 'w', encoding='utf-8', newline='') as csvfile: + csv_writer = csv.writer(csvfile, delimiter='\t') + + csv_writer.writerow(('GazeTimeStamp', + 'MediaTimeStamp', + 'MediaFrameIndex', + 'Gaze3dX', + 'Gaze3dY', + 'Gaze3dZ', + 'Gaze2dX', + 'Gaze2dY', + 'PupilDiaLeft', + 'PupilDiaRight', + 'Confidence')) + + for media_idx in range(*export_range): + media_timestamp = self.g_pool.timestamps[media_idx] + for g in self.g_pool.gaze_positions_by_frame[media_idx]: + try: + pupil_dia = {} + for p in g['base_data']: + pupil_dia[p['id']] = p['diameter_3d'] + + pixel_pos = denormalize(g['norm_pos'], self.g_pool.capture.frame_size, flip_y=True) + undistorted2d = self.g_pool.capture.intrinsics.undistortPoints(pixel_pos) + undistorted3d = np.ones(3) + undistorted3d[:2] = undistorted2d + undistorted2d = self.g_pool.capture.intrinsics.projectPoints(undistorted3d, use_distortion=False) + + data = (g['timestamp'], + media_timestamp, + media_idx - export_range[0], + *g['gaze_point_3d'], # Gaze3dX/Y/Z + *undistorted2d.flat, # Gaze2dX/Y + pupil_dia.get(1, 0.), # PupilDiaLeft + pupil_dia.get(0, 0.), # PupilDiaRight + g['confidence']) # Confidence + except KeyError: + if not user_warned_3d_only: + logger.error('Currently, the iMotions export only supports 3d gaze data') + user_warned_3d_only = True + continue + csv_writer.writerow(data) + + def recent_events(self, events): + if self.export_task: + recent = [d for d in self.export_task.fetch()] + if recent: + self.status, self.progress = recent[-1] + if self.export_task.canceled: + self.status = 'Export has been canceled' + self.progress = 0. + + def gl_display(self): + self.menu_icon.indicator_stop = self.progress / 100. + + def get_recording_start_date(self): + csv_loc = os.path.join(self.g_pool.rec_dir, 'info.csv') + with open(csv_loc, 'r', encoding='utf-8') as csvfile: + rec_info = csv_utils.read_key_value_file(csvfile) + date = rec_info['Start Date'].replace('.', '_').replace(':', '_') + time = rec_info['Start Time'].replace(':', '_') + return '{}_{}'.format(date, time)