aboutsummaryrefslogtreecommitdiff
path: root/src/argaze/utils/export_tobii_segment_plots.py
blob: 94e88e7afe6e3a1c1937e9745121aca42cc53db2 (plain)
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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
#!/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))

        # Plot pupil diameter data
        subplot = figure.add_subplot(711)
        subplot.set_title('Pupil diameter', loc='left')
        subplot.set_ylim(0, 10)

        patches = tobii_segment_data['PupilDiameter'].plot(names=['value'], colors=['#FFD800'])

        subplot.legend(handles=patches, loc='upper left')

        # Annotate events
        df_ts_events = tobii_segment_data['Event'].as_dataframe()

        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')
        subplot.set_ylim(-40, -20)

        patches = tobii_segment_data['PupilCenter'].plot(names=['x','y','z'], colors=['#276FB6','#9427B6','#888888'], split={'value':['x','y','z']})

        subplot.legend(handles=patches, 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')
        subplot.set_ylim(0., 1.)

        patches = tobii_segment_data['GazePosition'].plot(names=['x','y'], colors=['#276FB6','#9427B6'], split={'value':['x','y']})

        subplot.legend(handles=patches, loc='upper left')

        # Plot gaze direction data
        subplot = figure.add_subplot(714)
        subplot.set_title('Gaze direction', loc='left')

        patches = tobii_segment_data['GazeDirection'].plot(names=['x','y','z'], colors=['#276FB6','#9427B6','#888888'], split={'value':['x','y','z']})

        subplot.legend(handles=patches, loc='upper left')

        # Plot gaze direction data
        subplot = figure.add_subplot(715)
        subplot.set_title('Gaze position 3D', loc='left')

        patches = tobii_segment_data['GazePosition3D'].plot(names=['x','y','z'], colors=['#276FB6','#9427B6','#888888'], split={'value':['x','y','z']})

        subplot.legend(handles=patches, loc='upper left')

        # Plot accelerometer data
        subplot = figure.add_subplot(716)
        subplot.set_title('Accelerometer', loc='left')

        patches = tobii_segment_data['Accelerometer'].plot(names=['x','y','z'], colors=['#276FB6','#9427B6','#888888'], split={'value':['x','y','z']})

        subplot.legend(handles=patches, 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')

        patches = tobii_segment_data['Gyroscope'].plot(names=['x','y','z'], colors=['#276FB6','#9427B6','#888888'], split={'value':['x','y','z']})

        subplot.legend(handles=patches, 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()