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

#!/usr/bin/env python

import argparse
import os
import json

from argaze import DataStructures
from argaze import GazeFeatures
from argaze.TobiiGlassesPro2 import TobiiEntities, TobiiVideo, TobiiSpecifications
from argaze.ArUcoMarkers import *
from argaze.AreaOfInterest import *
from argaze.utils import MiscFeatures

import numpy
import cv2 as cv

def main():
    """
    Track ArUco markers into Tobii Glasses Pro 2 segment video file. 
    For each loaded AOI scene .obj file, position the scene virtually relatively to each detected ArUco markers and project the scene into camera frame. 
    Export AOIs video and data as a aruco_aoi.csv, aruco_aoi.mp4 files
    """

    # 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('-c', '--camera_calibration', metavar='CAM_CALIB', type=str, default=None, help='json camera calibration filepath')
    parser.add_argument('-p', '--aruco_tracker_configuration', metavar='TRACK_CONFIG', type=str, default=None, help='json aruco tracker configuration filepath')
    parser.add_argument('-md', '--marker_dictionary', metavar='MARKER_DICT', type=ArUcoMarkersDictionary.ArUcoMarkersDictionary, default='DICT_ARUCO_ORIGINAL', help='aruco marker dictionnary (DICT_4X4_50, DICT_4X4_100, DICT_4X4_250, DICT_4X4_1000, DICT_5X5_50, DICT_5X5_100, DICT_5X5_250, DICT_5X5_1000, DICT_6X6_50, DICT_6X6_100, DICT_6X6_250, DICT_6X6_1000, DICT_7X7_50, DICT_7X7_100, DICT_7X7_250, DICT_7X7_1000, DICT_ARUCO_ORIGINAL, DICT_APRILTAG_16h5, DICT_APRILTAG_25h9, DICT_APRILTAG_36h10, DICT_APRILTAG_36h11)')
    parser.add_argument('-ms', '--marker_size', metavar='MARKER_SIZE', type=float, default=6, help='aruco marker size (cm)')
    parser.add_argument('-mi', '--marker_id_scene', metavar='MARKER_ID_SCENE', type=json.loads, help='{"marker": "aoi scene filepath"} dictionary')
    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 markers id to track
        if args.marker_id_scene == None:
            print(f'Track any Aruco markers from the {args.marker_dictionary} dictionary')
        else:
            print(f'Track Aruco markers {list(args.marker_id_scene.keys())} from the {args.marker_dictionary} dictionary')

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

        vs_data_filepath = f'{destination_path}/aruco_aoi.csv'
        vs_video_filepath = f'{destination_path}/aruco_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'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 head rotations data buffer
        tobii_ts_head_rotations = tobii_segment_data['Gyroscope']

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

        # Create aruco camera
        aruco_camera = ArUcoCamera.ArUcoCamera()

        # Load calibration file
        if args.camera_calibration != None:

            aruco_camera.load_calibration_file(args.camera_calibration)

        else:

            raise UserWarning('.json camera calibration filepath required. Use -c option.')

        # Create aruco tracker
        aruco_tracker = ArUcoTracker.ArUcoTracker(args.marker_dictionary, args.marker_size, aruco_camera)

        # Load specific configuration file
        if args.aruco_tracker_configuration != None:

            aruco_tracker.load_configuration_file(args.aruco_tracker_configuration)

            print(f'ArUcoTracker configuration for {aruco_tracker.get_markers_dictionay().get_markers_format()} markers detection:')
            aruco_tracker.print_configuration()

        # Load AOI 3D scene for each marker and create a AOI 2D scene and frame when a 'Visualisation_Plan' AOI exist
        aoi3D_scenes = {}
        aoi2D_visu_scenes = {}

        for marker_id, aoi_scene_filepath in args.marker_id_scene.items():

            marker_id = int(marker_id)
            
            aoi3D_scenes[marker_id] = AOI3DScene.AOI3DScene()
            aoi3D_scenes[marker_id].load(aoi_scene_filepath)

            print(f'AOI in {os.path.basename(aoi_scene_filepath)} scene related to marker #{marker_id}:')
            for aoi in aoi3D_scenes[marker_id].keys():

                print(f'\t{aoi}')

        def aoi3D_scene_selector(marker_id):
            return aoi3D_scenes.get(marker_id, None)

        # Create timestamped buffer to store AOIs scene in time
        ts_aois_scenes = AOIFeatures.TimeStampedAOIScenes()

        # Video and data replay loop
        try:

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

            head_moving = False
            head_movement_last = 0.

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

                video_ts_ms = video_ts / 1e3

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

                # Process video and data frame
                try:

                    # Get nearest head rotation before video timestamp and remove all head rotations before
                    _, nearest_head_rotation = tobii_ts_head_rotations.pop_first_until(video_ts)

                    # Calculate head movement considering only head yaw and pitch
                    head_movement = numpy.array(nearest_head_rotation.value)
                    head_movement_px = head_movement.astype(int)
                    head_movement_norm = numpy.linalg.norm(head_movement[0:2])

                    # Draw movement vector
                    cv.line(visu_frame.matrix, (int(visu_frame.width/2), int(visu_frame.height/2)), (int(visu_frame.width/2) + head_movement_px[1], int(visu_frame.height/2) - head_movement_px[0]), (150, 150, 150), 3)
                
                    # Head movement detection hysteresis
                    # TODO : pass the threshold value as argument
                    if not head_moving and head_movement_norm > 50:
                        head_moving = True
                        
                    if head_moving and head_movement_norm < 10:
                        head_moving = False

                    # When head is moving, ArUco tracking could return bad pose estimation and so bad AOI scene projection
                    if head_moving:

                        ts_aois_scenes[round(video_ts_ms)] = AOIFeatures.EmptyAOIScene()
                        
                        raise UserWarning('Head is moving')

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

                    # Track markers with pose estimation and draw them
                    aruco_tracker.track(video_frame.matrix)
                    aruco_tracker.draw(visu_frame.matrix)

                    # When no marker is detected, no AOI scene projection can't be done
                    if aruco_tracker.markers_number() == 0:

                        ts_aois_scenes[round(video_ts_ms)] = AOIFeatures.EmptyAOIScene()

                        raise UserWarning('No marker detected')

                    # Store aoi 2D video for further scene merging
                    aoi2D_dict = {}

                    # Project 3D scene on each video frame and the visualisation frame
                    for (i, marker_id) in enumerate(aruco_tracker.get_markers_ids()):

                        # Copy 3D scene related to detected marker
                        aoi3D_scene = aoi3D_scene_selector(marker_id)
                        
                        if aoi3D_scene == None:
                            continue

                        # Transform scene into camera referential
                        aoi3D_camera = aoi3D_scene.transform(aruco_tracker.get_marker_translation(i), aruco_tracker.get_marker_rotation(i))

                        # Get aoi inside vision cone field 
                        cone_vision_height_cm = 200 # cm
                        cone_vision_radius_cm = numpy.tan(numpy.deg2rad(TobiiSpecifications.VISUAL_HFOV / 2)) * cone_vision_height_cm

                        aoi3D_inside, aoi3D_outside = aoi3D_camera.vision_cone(cone_vision_radius_cm, cone_vision_height_cm)

                        # Keep only aoi inside vision cone field
                        aoi3D_scene = aoi3D_scene.copy(exclude=aoi3D_outside.keys())

                        # DON'T APPLY CAMERA DISTORSION : it projects points which are far from the frame into it
                        # This hack isn't realistic but as the gaze will mainly focus on centered AOI, where the distorsion is low, it is acceptable.
                        aoi2D_video_scene = aoi3D_scene.project(aruco_tracker.get_marker_translation(i), aruco_tracker.get_marker_rotation(i), aruco_camera.get_K())

                        # Store each 2D aoi for further scene merging
                        for name, aoi in aoi2D_video_scene.items():

                            if name not in aoi2D_dict.keys():
                                aoi2D_dict[name] = []

                            aoi2D_dict[name].append(aoi.clockwise())

                    # Merge all 2D aoi into a single 2D scene
                    aoi2D_merged_scene = AOI2DScene.AOI2DScene()
                    for name, aoi_array in aoi2D_dict.items():
                        aoi2D_merged_scene[name] = numpy.sum(aoi_array, axis=0) / len(aoi_array)

                    aoi2D_merged_scene.draw(visu_frame.matrix, (0, 0))

                    # Store 2D merged scene at this time in millisecond
                    ts_aois_scenes[round(video_ts_ms)] = aoi2D_merged_scene

                    # Warn user when the merged scene is empty
                    if len(aoi2D_merged_scene.keys()) == 0:

                        raise UserWarning('Scene is empty')

                # Write warning
                except UserWarning as w:

                    cv.rectangle(visu_frame.matrix, (0, 50), (550, 100), (127, 127, 127), -1)
                    cv.putText(visu_frame.matrix, str(w), (20, 80), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv.LINE_AA)

                # Raised when timestamped buffer is empty
                except KeyError:
                    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'Segment time: {int(video_ts_ms)} 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.get_id()} ArUco AOI', visu_frame.matrix)

                # Write video
                output_video.write(visu_frame.matrix)

                # Update Progress Bar
                progress = video_ts_ms - int(args.time_range[0] * 1e3)
                MiscFeatures.printProgressBar(progress, tobii_segment_video.get_duration()/1e3, 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('\nAruco marker tracking metrics')
        try_count, tracked_counts, rejected_counts = aruco_tracker.get_track_metrics()

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

        for marker_id, rejected_count in rejected_counts.items():
            print(f'Markers {marker_id} has been rejected in {rejected_count} / {try_count} frames ({round(100 * rejected_count / try_count, 2)} %)')

        # Export aruco aoi data
        ts_aois_scenes.export_as_csv(vs_data_filepath, exclude=['dimension'])
        print(f'Aruco AOI data saved into {vs_data_filepath}')

        # Notify when the aruco aoi video has been exported
        print(f'Aruco AOI video saved into {vs_video_filepath}')

if __name__ == '__main__':

    main()