Adding a new argaze utils script to plot tobii segment data.

1 files changed, 210 insertions, 0 deletions

diff --git a/src/argaze/utils/export_tobii_segment_plots.py b/src/argaze/utils/export_tobii_segment_plots.py
new file mode 100644
index 0000000..13a6ead
--- /dev/null
+++ b/src/argaze/utils/export_tobii_segment_plots.py
@@ -0,0 +1,210 @@
+#!/usr/bin/env python
+
+import argparse
+import os
+import json
+
+from argaze import DataStructures
+from argaze import GazeFeatures
+from argaze.TobiiGlassesPro2 import TobiiEntities, TobiiVideo
+from argaze.utils import MiscFeatures
+
+import pandas
+import matplotlib.pyplot as mpyplot
+import matplotlib.patches as mpatches
+
+color = {
+    'Scene_Plan': (127, 127, 127),
+    'PFD_Plan': (63, 127, 63),
+    'Attitude_Plan': (0, 255, 0),
+    'Air_Speed_Plan': (255, 0, 255),
+    'Vertical_Speed_Plan': (255, 255, 0),
+    'Localiser_Plan': (0, 0, 255),
+    'ND_Plan': (127, 63, 63),
+    'Marker_Plan': (0, 0, 0)
+}
+
+def main():
+    """
+    """
+
+    # Manage arguments
+    parser = argparse.ArgumentParser(description=main.__doc__.split('-')[0])
+    parser.add_argument('-s', '--segment_path', metavar='SEGMENT_PATH', type=str, default=None, help='segment path')
+    parser.add_argument('-r', '--time_range', metavar=('START_TIME', 'END_TIME'), nargs=2, type=float, default=(0., None), help='start and end time (in second)')
+    
+    parser.add_argument('-o', '--output', metavar='OUT', type=str, default=None, help='destination folder path (segment folder by default)')
+    args = parser.parse_args()
+
+    if args.segment_path != None:
+
+        # Manage destination path
+        destination_path = '.'
+        if args.output != None:
+
+            if not os.path.exists(os.path.dirname(args.output)):
+
+                os.makedirs(os.path.dirname(args.output))
+                print(f'{os.path.dirname(args.output)} folder created')
+
+                destination_path = args.output
+
+        else:
+
+                destination_path = args.segment_path
+
+        data_plots_filepath = f'{destination_path}/plots.svg'
+
+        # Load a tobii segment
+        tobii_segment = TobiiEntities.TobiiSegment(args.segment_path, int(args.time_range[0] * 1e6), int(args.time_range[1] * 1e6) if args.time_range[1] != None else None)
+
+        # Load a tobii segment video
+        tobii_segment_video = tobii_segment.load_video()
+        print(f'Video properties:\n\tduration: {tobii_segment_video.get_duration() / 1e6} s\n\twidth: {tobii_segment_video.get_width()} px\n\theight: {tobii_segment_video.get_height()} px')
+
+        # Load a tobii segment data
+        tobii_segment_data = tobii_segment.load_data()
+        
+        print(f'Loaded data count:')
+        for name in tobii_segment_data.keys():
+            print(f'\t{name}: {len(tobii_segment_data[name])} data')
+
+        # Edit figure
+        figure = mpyplot.figure(figsize=(4 * tobii_segment_video.get_duration() / 1e6, 35))
+
+        # Access to timestamped pupil diameter data buffer
+        df_ts_pupil_diameters = tobii_segment_data['PupilDiameter'].as_dataframe()
+
+        # Plot pupil diameter data
+        subplot = figure.add_subplot(711)
+        subplot.set_title('Pupil diameter', loc='left')
+
+        mpyplot.ylim(0, 10)
+
+        MiscFeatures.plotTimestampedDataframe(df_ts_pupil_diameters, 'value', color='#FFD800')
+
+        value_patch = mpatches.Patch(color='#FFD800', label='DIAMETER')
+        subplot.legend(handles=[value_patch], loc='upper left')
+
+        # Access to timestamped event data buffer
+        df_ts_events = tobii_segment_data['Event'].as_dataframe()
+
+        # Annotate events
+        if len(df_ts_events) > 0:
+            
+            for ts, event_type, event_tag in zip(df_ts_events.index, df_ts_events.type, df_ts_events.tag):
+                subplot.annotate(f'{event_type}\n{event_tag}', xy=(ts, 7), horizontalalignment="left", verticalalignment="top")
+                subplot.vlines(ts, 0, 6, color="tab:red", linewidth=1)
+
+        # Access to timestamped pupil center data buffer
+        df_ts_pupil_centers = tobii_segment_data['PupilCenter'].as_dataframe(split={'value':['x','y', 'z']})
+
+        # Plot pupil center data
+        subplot = figure.add_subplot(712)
+        subplot.set_title('Pupil center', loc='left')
+
+        mpyplot.ylim(-40, -20)
+
+        MiscFeatures.plotTimestampedDataframe(df_ts_pupil_centers, 'x', color='#276FB6')
+        MiscFeatures.plotTimestampedDataframe(df_ts_pupil_centers, 'y', color='#9427B6')
+        MiscFeatures.plotTimestampedDataframe(df_ts_pupil_centers, 'z', color='#888888')
+
+        x_patch = mpatches.Patch(color='#276FB6', label='X')
+        y_patch = mpatches.Patch(color='#9427B6', label='Y')
+        z_patch = mpatches.Patch(color='#888888', label='Z')
+        subplot.legend(handles=[x_patch, y_patch, z_patch], loc='upper left')
+
+        # Access to timestamped gaze position data buffer
+        df_ts_gaze_positions = tobii_segment_data['GazePosition'].as_dataframe(split={'value':['x','y']})
+
+        # Plot gaze position data
+        subplot = figure.add_subplot(713)
+        subplot.set_title('Gaze position', loc='left')
+
+        mpyplot.ylim(0., 1.)
+
+        MiscFeatures.plotTimestampedDataframe(df_ts_gaze_positions, 'x', color='#276FB6')
+        MiscFeatures.plotTimestampedDataframe(df_ts_gaze_positions, 'y', color='#9427B6')
+
+        x_patch = mpatches.Patch(color='#276FB6', label='X')
+        y_patch = mpatches.Patch(color='#9427B6', label='Y')
+        subplot.legend(handles=[x_patch, y_patch], loc='upper left')
+
+        # Access to timestamped gaze direction data buffer
+        df_ts_gaze_directions = tobii_segment_data['GazeDirection'].as_dataframe(split={'value':['x','y', 'z']})
+
+        # Plot gaze direction data
+        subplot = figure.add_subplot(714)
+        subplot.set_title('Gaze direction', loc='left')
+
+        #mpyplot.ylim(0., 1.)
+
+        MiscFeatures.plotTimestampedDataframe(df_ts_gaze_directions, 'x', color='#276FB6')
+        MiscFeatures.plotTimestampedDataframe(df_ts_gaze_directions, 'y', color='#9427B6')
+        MiscFeatures.plotTimestampedDataframe(df_ts_gaze_directions, 'z', color='#888888')
+
+        x_patch = mpatches.Patch(color='#276FB6', label='X')
+        y_patch = mpatches.Patch(color='#9427B6', label='Y')
+        z_patch = mpatches.Patch(color='#888888', label='Z')
+        subplot.legend(handles=[x_patch, y_patch, z_patch], loc='upper left')
+
+        # Access to timestamped gaze position 3d data buffer
+        df_ts_gaze_positions_3d = tobii_segment_data['GazePosition3D'].as_dataframe(split={'value':['x','y', 'z']})
+
+        # Plot gaze direction data
+        subplot = figure.add_subplot(715)
+        subplot.set_title('Gaze position 3D', loc='left')
+
+        #mpyplot.ylim(0., 1.)
+
+        MiscFeatures.plotTimestampedDataframe(df_ts_gaze_positions_3d, 'x', color='#276FB6')
+        MiscFeatures.plotTimestampedDataframe(df_ts_gaze_positions_3d, 'y', color='#9427B6')
+        MiscFeatures.plotTimestampedDataframe(df_ts_gaze_positions_3d, 'z', color='#888888')
+
+        x_patch = mpatches.Patch(color='#276FB6', label='X')
+        y_patch = mpatches.Patch(color='#9427B6', label='Y')
+        z_patch = mpatches.Patch(color='#888888', label='Z')
+        subplot.legend(handles=[x_patch, y_patch, z_patch], loc='upper left')
+
+        # Access to timestamped accelerometer data buffer
+        df_ts_accelerometers = tobii_segment_data['Accelerometer'].as_dataframe(split={'value':['x','y', 'z']})
+
+        # Plot accelerometer data
+        subplot = figure.add_subplot(716)
+        subplot.set_title('Accelerometer', loc='left')
+
+        MiscFeatures.plotTimestampedDataframe(df_ts_accelerometers, 'x', color='#276FB6')
+        MiscFeatures.plotTimestampedDataframe(df_ts_accelerometers, 'y', color='#9427B6')
+        MiscFeatures.plotTimestampedDataframe(df_ts_accelerometers, 'z', color='#888888')
+
+        x_patch = mpatches.Patch(color='#276FB6', label='X')
+        y_patch = mpatches.Patch(color='#9427B6', label='Y')
+        z_patch = mpatches.Patch(color='#888888', label='Z')
+        subplot.legend(handles=[x_patch, y_patch, z_patch], loc='upper left')
+
+        # Access to timestamped gyroscope data buffer
+        df_ts_gyroscopes = tobii_segment_data['Gyroscope'].as_dataframe(split={'value':['x','y', 'z']})
+
+        # Plot accelerometer data
+        subplot = figure.add_subplot(717)
+        subplot.set_title('Gyroscope', loc='left')
+
+        MiscFeatures.plotTimestampedDataframe(df_ts_gyroscopes, 'x', color='#276FB6')
+        MiscFeatures.plotTimestampedDataframe(df_ts_gyroscopes, 'y', color='#9427B6')
+        MiscFeatures.plotTimestampedDataframe(df_ts_gyroscopes, 'z', color='#888888')
+
+        x_patch = mpatches.Patch(color='#276FB6', label='X')
+        y_patch = mpatches.Patch(color='#9427B6', label='Y')
+        z_patch = mpatches.Patch(color='#888888', label='Z')
+        subplot.legend(handles=[x_patch, y_patch, z_patch], loc='upper left')
+
+        # Export data plots
+        mpyplot.tight_layout()
+        mpyplot.savefig(data_plots_filepath)
+        mpyplot.close('all')
+
+        print(f'\nData plots saved into {data_plots_filepath}')
+
+if __name__ == '__main__':
+
+    main()