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

#!/usr/bin/env python

import argparse
import os, json
import math
import threading

from argaze import *
from argaze.TobiiGlassesPro2 import *
from argaze.ArUcoMarkers import *
from argaze.AreaOfInterest import *
from argaze.utils import MiscFeatures

import cv2 as cv
import numpy

def main():
    """
    Track ArUcoPlan into Tobii Glasses Pro 2 camera video stream.
    """

    # 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('-t', '--time_range', metavar=('START_TIME', 'END_TIME'), nargs=2, type=float, default=(0., None), help='start and end time (in second)')
    parser.add_argument('-p', '--project_path', metavar='ARGAZE_PROJECT', type=str, default=None, help='json argaze project filepath')
    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')

        aoi_json_filepath = f'{destination_path}/aoi.json'
        aoi_csv_filepath = f'{destination_path}/aoi.csv'
        aoi_mp4_filepath = f'{destination_path}/aoi.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'\nVideo 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'\nLoaded data count:')
        for name in tobii_segment_data.keys():
            print(f'\t{name}: {len(tobii_segment_data[name])} data')

        # Access to video timestamp data buffer
        tobii_ts_vts = tobii_segment_data['VideoTimeStamp']

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

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

        # Initialise progress bar
        MiscFeatures.printProgressBar(0, tobii_segment_video.duration, prefix = '\nGazePositions projection:', suffix = 'Complete', length = 100)

        for ts, tobii_gaze_position in tobii_ts_gaze_positions.items():

            # Update Progress Bar
            progress = ts - int(args.time_range[0] * 1e6)
            MiscFeatures.printProgressBar(progress, tobii_segment_video.duration, prefix = 'GazePositions projection:', suffix = 'Complete', length = 100)


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

                gaze_position_px = (int(tobii_gaze_position.value[0] * tobii_segment_video.width), int(tobii_gaze_position.value[1] * tobii_segment_video.height))
                ts_gaze_positions[ts] = GazeFeatures.GazePosition(gaze_position_px) 

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

        # Load ar scene
        ar_scene = ArScene.ArScene.from_json(args.project_path)

        print(ar_scene)

        # Create timestamped buffer to store AOIs and primary time stamp offset
        ts_offset_aois = DataStructures.TimeStampedBuffer()

        # Video and data replay loop
        try:

            # Initialise progress bar
            MiscFeatures.printProgressBar(0, tobii_segment_video.duration/1e3, prefix = 'Progress:', suffix = 'Complete', length = 100)

            # Iterate on video frames
            for video_ts, video_frame in tobii_segment_video.frames():

                # This video frame is the reference until the next frame
                # Here next frame is at + 40ms (25 fps)
                # TODO: Get video fps to adapt
                next_video_ts = video_ts + 40000 

                # Copy video frame to edit visualisation on it without disrupting aruco tracking
                visu_frame = video_frame.copy()

                # Prepare to store projected AOI
                projected_aois = {}

                # Process video and data frame
                try:

                    # Get nearest video timestamp
                    _, nearest_vts = tobii_ts_vts.get_last_before(video_ts)

                    projected_aois['offset'] = nearest_vts.offset

                    # Hide frame left and right borders before tracking to ignore markers outside focus area
                    cv.rectangle(video_frame.matrix, (0, 0), (int(video_frame.width/6), int(video_frame.height)), (0, 0, 0), -1)
                    cv.rectangle(video_frame.matrix, (int(video_frame.width*(1 - 1/6)), 0), (int(video_frame.width), int(video_frame.height)), (0, 0, 0), -1)

                    # Project scene into frame
                    scene_projection = ar_scene.project(video_frame.matrix, consistent_markers_number=1, visual_hfov=TobiiSpecifications.VISUAL_HFOV)

                    # Store all projected aoi
                    for aoi_name in scene_projection.keys():

                        projected_aois[aoi_name] = numpy.rint(scene_projection[aoi_name]).astype(int)

                    # Draw tracked markers
                    ar_scene.aruco_tracker.draw_tracked_markers(visu_frame.matrix)

                    # Draw scene projection
                    scene_projection.draw(visu_frame.matrix, (0, 0), color=(0, 255, 255))

                # Catch exceptions raised by project_scene method
                except (ArScene.PoseEstimationFailed, ArScene.SceneProjectionFailed) as e:

                    # Draw tracked markers
                    ar_scene.aruco_tracker.draw_tracked_markers(visu_frame.matrix)

                    if str(e) == 'Unconsistent marker poses':

                        projected_aois['error'] = str(e) + ': ' + str(e.unconsistencies)

                    else:

                        projected_aois['error'] = str(e)

                    cv.rectangle(visu_frame.matrix, (0, 100), (550, 150), (127, 127, 127), -1)
                    cv.putText(visu_frame.matrix, str(e), (20, 130), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv.LINE_AA)

                # Raised when timestamped buffer is empty
                except KeyError as e:

                    e = 'VideoTimeStamp missing'

                    projected_aois['offset'] = 0
                    projected_aois['error'] = e

                    cv.rectangle(visu_frame.matrix, (0, 100), (550, 150), (127, 127, 127), -1)
                    cv.putText(visu_frame.matrix, str(e), (20, 130), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv.LINE_AA)

                # Store projected AOI
                ts_offset_aois[video_ts] = projected_aois

                # Draw gaze positions until next frame
                try:

                    # Get next gaze position
                    ts_start, start_gaze_position = ts_gaze_positions.first
                    ts_next, next_gaze_position = ts_gaze_positions.first
                    
                    # Check next gaze position is not after next frame time
                    while ts_next < next_video_ts:

                        ts_start, start_gaze_position = ts_gaze_positions.pop_first()
                        ts_next, next_gaze_position = ts_gaze_positions.first

                        # Draw start gaze
                        start_gaze_position.draw(visu_frame.matrix)

                        if start_gaze_position.valid and next_gaze_position.valid:

                            # Draw movement from start to next
                            cv.line(visu_frame.matrix, start_gaze_position, next_gaze_position, (0, 255, 255), 1)

                    if start_gaze_position.valid:

                        # Write last start gaze position
                        cv.putText(visu_frame.matrix, str(start_gaze_position.value), start_gaze_position.value, cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv.LINE_AA)
                        
                    # Write last start gaze position timing
                    cv.rectangle(visu_frame.matrix, (0, 50), (550, 100), (31, 31, 31), -1)
                    cv.putText(visu_frame.matrix, f'Gaze time: {ts_start*1e-3:.3f} ms', (20, 85), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv.LINE_AA)
        
                # Empty gaze position
                except IndexError:
                    pass

                # Draw focus area
                cv.rectangle(visu_frame.matrix, (int(video_frame.width/6), 0), (int(visu_frame.width*(1-1/6)), int(visu_frame.height)), (255, 150, 150), 1)
                       
                # Draw center
                cv.line(visu_frame.matrix, (int(visu_frame.width/2) - 50, int(visu_frame.height/2)), (int(visu_frame.width/2) + 50, int(visu_frame.height/2)), (255, 150, 150), 1)
                cv.line(visu_frame.matrix, (int(visu_frame.width/2), int(visu_frame.height/2) - 50), (int(visu_frame.width/2), int(visu_frame.height/2) + 50), (255, 150, 150), 1)

                # Write segment timing
                cv.rectangle(visu_frame.matrix, (0, 0), (550, 50), (63, 63, 63), -1)
                cv.putText(visu_frame.matrix, f'Video time: {video_ts*1e-3:.3f} ms', (20, 40), 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 visualisation
                    cv.imshow(f'Segment {tobii_segment.id} ArUco AOI', visu_frame.matrix)

                # Write video
                output_video.write(visu_frame.matrix)

                # Update Progress Bar
                progress = video_ts*1e-3 - int(args.time_range[0] * 1e3)
                MiscFeatures.printProgressBar(progress, tobii_segment_video.duration*1e-3, prefix = 'Progress:', suffix = 'Complete', length = 100)

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

        # Stop frame display
        cv.destroyAllWindows()

        # End output video file
        output_video.close()

        # Print aruco tracking metrics
        print('\n\nAruco marker tracking metrics')
        try_count, tracked_counts = ar_scene.aruco_tracker.track_metrics

        for marker_id, tracked_count in tracked_counts.items():
            print(f'\tMarkers {marker_id} has been detected in {tracked_count} / {try_count} frames ({round(100 * tracked_count / try_count, 2)} %)')

        # Export aruco aoi data
        ts_offset_aois.to_json(aoi_json_filepath)
        ts_offset_aois.as_dataframe().to_csv(aoi_csv_filepath)
        print(f'Aruco AOI data saved into {aoi_json_filepath} and {aoi_csv_filepath}')

        # Notify when the aruco aoi video has been exported
        print(f'Aruco AOI video saved into {aoi_mp4_filepath}')
    
if __name__ == '__main__':

    main()