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

#!/usr/bin/env python

import argparse
import os

from argaze import GazeFeatures
from argaze.TobiiGlassesPro2 import TobiiEntities, TobiiVideo, TobiiSpecifications
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=50, 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
                if args.time_range[1] != None:
                    timerange_path = f'[{int(args.time_range[0])}s - {int(args.time_range[1])}s]'
                else:
                    timerange_path = f'[all]'

                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.duration/1e6} s\n\twidth: {tobii_segment_video.width} px\n\theight: {tobii_segment_video.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']

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

        # Format tobii gaze position and accuracy in pixel
        ts_gaze_positions = GazeFeatures.TimeStampedGazePositions()

        for ts, tobii_gaze_position in tobii_ts_gaze_positions.items():

            # Test gaze position validity
            if tobii_gaze_position.validity == 0:

                gaze_position_pixel = GazeFeatures.GazePosition( (int(tobii_gaze_position.value[0] * tobii_segment_video.width), int(tobii_gaze_position.value[1] * tobii_segment_video.height)) )
                
                # Get gaze position 3D at same gaze position timestamp
                tobii_gaze_position_3d = tobii_ts_gaze_positions_3d.pop(ts)

                # Test gaze position 3d validity
                if tobii_gaze_position_3d.validity == 0:
                    
                    gaze_accuracy_mm = numpy.sin(numpy.deg2rad(TobiiSpecifications.ACCURACY)) * tobii_gaze_position_3d.value[2]
                    tobii_camera_hfov_mm = numpy.sin(numpy.deg2rad(TobiiSpecifications.CAMERA_HFOV)) * tobii_gaze_position_3d.value[2]
                    
                    gaze_position_pixel.accuracy = round(tobii_segment_video.width * float(gaze_accuracy_mm) / float(tobii_camera_hfov_mm))

                    # Store gaze position using millisecond unit timestamp
                    ts_gaze_positions[ts/1e3] = gaze_position_pixel

                    continue

            # Store unvalid gaze position for further movement processing
            ts_gaze_positions[ts/1e3] = GazeFeatures.UnvalidGazePosition()

        print(f'Movement identifier parameters:')
        print(f'\tDispersion threshold = {args.dispersion_threshold}')
        print(f'\tDuration 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.duration/1e3), 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] * 1e3)
            MiscFeatures.printProgressBar(progress, int(tobii_segment_video.duration/1e3), 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.stream)

        # Video and data loop
        try:

            # Initialise progress bar
            MiscFeatures.printProgressBar(0, tobii_segment_video.duration/1e3, 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 / 1e3

                # 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 = 1

                        # 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, (0, 255, 0), current_fixation_time_counter)

                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)

                    # Draw gaze
                    nearest_gaze_position.draw(video_frame.matrix)

                # Wait for gaze position
                except KeyError:
                    pass

                # Write segment timing
                cv.rectangle(video_frame.matrix, (0, 0), (550, 50), (63, 63, 63), -1)
                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.rectangle(video_frame.matrix, (0, 90), (550, 150), (63, 63, 63), -1)
                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)

                if args.window:

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

                    # Display video
                    cv.imshow(f'Segment {tobii_segment.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] * 1e3)
                MiscFeatures.printProgressBar(progress, tobii_segment_video.duration/1e3, 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()