path: root/src/argaze/utils/export_tobii_segment_movements.py

#!/usr/bin/env python

import argparse
import os

from argaze import GazeFeatures
from argaze.TobiiGlassesPro2 import TobiiEntities, TobiiVideo
from argaze.utils import MiscFeatures

import cv2 as cv
import numpy

def main():
    """
    Analyse Tobii segment fixations
    """

    # Manage arguments
    parser = argparse.ArgumentParser(description=main.__doc__.split('-')[0])
    parser.add_argument('-s', '--segment_path', metavar='SEGMENT_PATH', type=str, default=None, help='path to a tobii segment folder')
    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('-d', '--dispersion_threshold', metavar='DISPERSION_THRESHOLD', type=int, default=10, help='dispersion threshold in pixel')
    parser.add_argument('-t', '--duration_threshold', metavar='DURATION_THRESHOLD', type=int, default=100, help='duration threshold in millisecond')
    parser.add_argument('-o', '--output', metavar='OUT', type=str, default=None, help='destination folder path (segment folder by default)')
    parser.add_argument('-w', '--window', metavar='DISPLAY', type=bool, default=True, help='enable window display', action=argparse.BooleanOptionalAction)
    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

                # Export into a dedicated time range folder
                timerange_path = f'[{int(args.time_range[0])}s - {int(args.time_range[1])}s]'

                destination_path = f'{destination_path}/{timerange_path}'

                if not os.path.exists(destination_path):

                    os.makedirs(destination_path)
                    print(f'{destination_path} folder created')

        fixations_filepath = f'{destination_path}/movements_fixations.csv'
        saccades_filepath = f'{destination_path}/movements_saccades.csv'

        gaze_status_filepath = f'{destination_path}/gaze_status.csv'
        gaze_status_video_filepath = f'{destination_path}/gaze_status.mp4'

        # 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')

        # Access to timestamped gaze position data buffer
        tobii_ts_gaze_positions = tobii_segment_data['GazePosition']

        # Format tobii gaze position in pixel and store them using millisecond unit timestamp
        ts_gaze_positions = GazeFeatures.TimeStampedGazePositions()

        for ts, tobii_gaze_position in tobii_ts_gaze_positions.items():
            gaze_position_pixel = (int(tobii_gaze_position.value[0] * tobii_segment_video.get_width()), int(tobii_gaze_position.value[1] * tobii_segment_video.get_height()))
            ts_gaze_positions[ts/1000] = gaze_position_pixel

        # Access to timestamped gaze 3D positions data buffer
        tobii_ts_gaze_positions_3d = tobii_segment_data['GazePosition3D']

        # Format gaze accuracies in pixel and store them using millisecond unit timestamp
        ts_gaze_accuracies = GazeFeatures.TimeStampedGazeAccuracies()

        # !!! the parameters below are specific to the TobiiGlassesPro2 !!!
        # Reference : https://www.biorxiv.org/content/10.1101/299925v1
        tobii_accuracy = 1.42 # degree
        tobii_precision = 0.34 # degree
        tobii_camera_hfov = 82 # degree

        for ts, tobii_ts_gaze_position_3d in tobii_ts_gaze_positions_3d.items():

            if tobii_ts_gaze_position_3d.value[2] > 0:

                gaze_accuracy_mm = numpy.sin(numpy.deg2rad(tobii_accuracy)) * tobii_ts_gaze_position_3d.value[2]
                tobii_camera_hfov_mm = numpy.sin(numpy.deg2rad(tobii_camera_hfov)) * tobii_ts_gaze_position_3d.value[2]
                gaze_accuracy_pixel = round(tobii_segment_video.get_width() * float(gaze_accuracy_mm) / float(tobii_camera_hfov_mm))

                ts_gaze_accuracies[ts/1000] = gaze_accuracy_pixel

        print(f'Dispersion threshold: {args.dispersion_threshold}')
        print(f'Duration threshold: {args.duration_threshold}')

        # Start movement identification
        movement_identifier = GazeFeatures.DispersionBasedMovementIdentifier(ts_gaze_positions, args.dispersion_threshold, args.duration_threshold)
        fixations = GazeFeatures.TimeStampedMovements()
        saccades = GazeFeatures.TimeStampedMovements()
        gaze_status = GazeFeatures.TimeStampedGazeStatus()

        # Initialise progress bar
        MiscFeatures.printProgressBar(0, int(tobii_segment_video.get_duration()/1000), prefix = 'Movements identification:', suffix = 'Complete', length = 100)

        for item in movement_identifier:

            if isinstance(item, GazeFeatures.DispersionBasedMovementIdentifier.DispersionBasedFixation):

                start_ts, start_position = item.positions.get_first()

                fixations[start_ts] = item

                for ts, position in item.positions.items():

                    gaze_status[ts] = GazeFeatures.GazeStatus(position, 'Fixation', len(fixations))

            elif isinstance(item, GazeFeatures.DispersionBasedMovementIdentifier.DispersionBasedSaccade):

                start_ts, start_position = item.positions.get_first()
                end_ts, end_position = item.positions.get_last()
                
                saccades[start_ts] = item

                gaze_status[start_ts] = GazeFeatures.GazeStatus(start_position, 'Saccade', len(saccades))
                gaze_status[end_ts] = GazeFeatures.GazeStatus(end_position, 'Saccade', len(saccades))

            else:
                continue

            # Update Progress Bar
            progress = ts - int(args.time_range[0] * 1000)
            MiscFeatures.printProgressBar(progress, int(tobii_segment_video.get_duration()/1000), prefix = 'Movements identification:', suffix = 'Complete', length = 100)

        print(f'\n{len(fixations)} fixations and {len(saccades)} saccades found')

        # Export fixations analysis
        fixations.export_as_csv(fixations_filepath)
        print(f'Fixations saved into {fixations_filepath}')

        # Export saccades analysis
        saccades.export_as_csv(saccades_filepath)
        print(f'Saccades saved into {saccades_filepath}')

        # Export gaze status analysis
        gaze_status.export_as_csv(gaze_status_filepath)
        print(f'Gaze status saved into {gaze_status_filepath}')

        # Prepare video exportation at the same format than segment video
        output_video = TobiiVideo.TobiiVideoOutput(gaze_status_video_filepath, tobii_segment_video.get_stream())

        # Video and data loop
        try:

            # Initialise progress bar
            MiscFeatures.printProgressBar(0, tobii_segment_video.get_duration()/1000, prefix = 'Video with movements processing:', suffix = 'Complete', length = 100)

            current_fixation_ts, current_fixation = fixations.pop_first()
            current_fixation_time_counter = 0

            current_saccade_ts, current_saccade = saccades.pop_first()
            
            # Iterate on video frames
            for video_ts, video_frame in tobii_segment_video.frames():

                video_ts_ms = video_ts / 1000

                # write segment timing
                cv.putText(video_frame.matrix, f'Segment time: {int(video_ts_ms)} ms', (20, 40), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv.LINE_AA)

                # write movement identification parameters
                cv.putText(video_frame.matrix, f'Dispersion threshold: {args.dispersion_threshold} px', (20, 100), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv.LINE_AA)
                cv.putText(video_frame.matrix, f'Duration threshold: {args.duration_threshold} ms', (20, 140), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv.LINE_AA)

                # Draw current fixation
                if len(fixations) > 0:

                    if video_ts_ms > current_fixation_ts + current_fixation.duration:

                        current_fixation_ts, current_fixation = fixations.pop_first()
                        current_fixation_time_counter = 0

                        # Draw saccade
                        if len(saccades) > 0:

                            if video_ts_ms > current_saccade_ts + current_saccade.duration:

                                current_saccade_ts, current_saccade = saccades.pop_first()
                                start_ts, start_position = current_saccade.positions.pop_first()
                                end_ts, end_position = current_saccade.positions.pop_first()

                            cv.line(video_frame.matrix, start_position, end_position, (0, 0, 255), 2)

                    else:

                        current_fixation_time_counter += 1

                    cv.circle(video_frame.matrix, current_fixation.centroid, current_fixation.dispersion + current_fixation_time_counter, (0, 255, 0), 1)

                try:

                    # Get closest gaze position before video timestamp and remove all gaze positions before
                    _, nearest_gaze_position = ts_gaze_positions.pop_first_until(video_ts_ms)

                    # Get closest gaze accuracy before video timestamp and remove all gaze accuracies before
                    _, nearest_gaze_accuracy = ts_gaze_accuracies.pop_first_until(video_ts_ms)

                    # Draw gaze position and precision    
                    cv.circle(video_frame.matrix, nearest_gaze_position, 2, (0, 255, 255), -1)
                    cv.circle(video_frame.matrix, nearest_gaze_position, nearest_gaze_accuracy, (0, 255, 255), 1)

                # Wait for gaze position
                except ValueError:
                    pass

                if args.window:

                    # Close window using 'Esc' key
                    if cv.waitKey(1) == 27:
                        break

                    # Display video
                    cv.imshow(f'Segment {tobii_segment.get_id()} movements', video_frame.matrix)

                # Write video
                output_video.write(video_frame.matrix)

                # Update Progress Bar
                progress = video_ts_ms - int(args.time_range[0] * 1000)
                MiscFeatures.printProgressBar(progress, tobii_segment_video.get_duration()/1000, prefix = 'Video with movements processing:', suffix = 'Complete', length = 100)

        # Exit on 'ctrl+C' interruption
        except KeyboardInterrupt:
            pass

        # End output video file
        output_video.close()
        print(f'\nVideo with movements saved into {gaze_status_video_filepath}')

if __name__ == '__main__':

    main()