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

#!/usr/bin/env python

import argparse
import os
import json
import time

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

import numpy
import cv2 as cv

def main():
    """
    Open video file with ArUco marker scene inside
    """

    # 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=str, 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
                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')

        #vs_data_filepath = f'{destination_path}/visual_scan.csv'

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

        # 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 ValueError('.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
        def load_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_configuration_file()

        # Load AOI 3D scene for each marker
        aoi3D_scenes = {}
        aoi3D_scene_edits = {}
    
        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)

            aoi3D_scene_edits[marker_id] = {
                'rotation': numpy.array([0.0, 0.0, 0.0]),
                'translation': numpy.array([0.0, 0.0, 0.0])
            }

            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)

        def aoi3D_scene_edit_selector(marker_id):
            return aoi3D_scene_edits.get(marker_id, None)

        # Display first frame
        video_ts, video_frame = tobii_segment_video.get_frame(0)
        cv.imshow(f'Segment {tobii_segment.get_id()} ArUco marker editor', video_frame.matrix)

        # Init mouse interaction variables
        pointer = (0, 0)
        left_click = (0, 0)
        right_click = (0, 0)
        right_button = False
        edit_coord = 0 # x

        # On mouse left left_click : update pointer position
        def on_mouse_event(event, x, y, flags, param):

            nonlocal pointer
            nonlocal left_click
            nonlocal right_click
            nonlocal right_button

            # Update pointer
            pointer = (x, y)

            # Update left_click
            if event == cv.EVENT_LBUTTONUP:

                left_click = pointer

            # Udpate right_button
            elif event == cv.EVENT_RBUTTONDOWN:

                right_button = True

            elif event == cv.EVENT_RBUTTONUP:

                right_button = False

            # Udpate right_click
            if right_button:

                right_click = pointer

        cv.setMouseCallback(f'Segment {tobii_segment.get_id()} ArUco marker editor', on_mouse_event)

        # Frame selector loop
        frame_index = 0
        last_frame_index = -1
        last_frame_matrix = video_frame.matrix.copy()

        selected_marker_id = -1

        try:

            while True:

                # Select a frame on change
                if frame_index != last_frame_index:

                    video_ts, video_frame = tobii_segment_video.get_frame(frame_index)
                    video_ts_ms = video_ts / 1000

                    last_frame_index = frame_index
                    last_frame_matrix = video_frame.matrix.copy()

                else:

                    video_frame.matrix = last_frame_matrix.copy()

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

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

                # Draw focus area
                cv.circle(video_frame.matrix, (int(video_frame.width/2), int(video_frame.height/2)), int(video_frame.width/3), (255, 150, 150), 1)

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

                # Write selected marker id
                if selected_marker_id >= 0:

                    cv.putText(video_frame.matrix, f'Marker {selected_marker_id} : Axis {edit_coord + 1} selected', (20, 80), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv.LINE_AA)

                # Write documentation
                else: 
                    cv.putText(video_frame.matrix, f'Left click on marker to select scene', (20, 80), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv.LINE_AA)
                    cv.putText(video_frame.matrix, f'Shift+num to select axis', (20, 120), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv.LINE_AA)
                    cv.putText(video_frame.matrix, f'Right click and drag to edit axis', (20, 160), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv.LINE_AA)
                    cv.putText(video_frame.matrix, f'Ctrl+s to save scene', (20, 200), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv.LINE_AA)

                # Project 3D scene on each video frame and the visualisation frame
                if aruco_tracker.get_markers_number():

                    # Write detected marker ids
                    cv.putText(video_frame.matrix, f'Detected markers : {aruco_tracker.get_markers_ids()}', (20, video_frame.height - 40), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv.LINE_AA)
                            
                    # Update selected marker id by left_clicking on marker
                    for (i, marker_id) in enumerate(aruco_tracker.get_markers_ids()):

                        marker_aoi = numpy.array(aruco_tracker.get_marker_corners(i)).view(AOIFeatures.AreaOfInterest)

                        if marker_aoi.looked(left_click):

                            selected_marker_id = marker_id

                            # Select 3D scene related to selected marker
                            aoi3D_scene = aoi3D_scene_selector(selected_marker_id)

                    # If a marker is selected
                    try:
                        
                        # Retreive marker index
                        selected_marker_index = aruco_tracker.get_marker_index(selected_marker_id)

                        # If AOI scene is found
                        if aoi3D_scene != None:

                            # Is the marker out of focus area ?
                            marker_x, marker_y = aruco_tracker.get_marker_center(selected_marker_index)
                            distance_to_center = ( (video_frame.width/2 - marker_x)**2 + (video_frame.height/2 - marker_y)**2 )**0.5

                            if distance_to_center > int(video_frame.width/3):

                                # Write warning
                                cv.putText(video_frame.matrix, f'Out of focus area', (20, 120), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv.LINE_AA)

                            # Select scene edit
                            aoi3D_scene_edit = aoi3D_scene_edit_selector(selected_marker_id)

                            # Edit scene
                            if aoi3D_scene_edit != None:

                                if right_button:

                                    pointer_delta_x, pointer_delta_y = (right_click[0] - video_frame.width/2) / (video_frame.width/3), (video_frame.height/2 - right_click[1]) / (video_frame.width/3)
                                    
                                    # Edit scene rotation
                                    if edit_coord == 0:
                                        aoi3D_scene_edit['rotation'] = numpy.array([pointer_delta_y, aoi3D_scene_edit['rotation'][1], aoi3D_scene_edit['rotation'][2]])

                                    elif edit_coord == 1:
                                        aoi3D_scene_edit['rotation'] = numpy.array([aoi3D_scene_edit['rotation'][0], pointer_delta_x, aoi3D_scene_edit['rotation'][2]])

                                    elif edit_coord == 2:
                                        aoi3D_scene_edit['rotation'] = numpy.array([aoi3D_scene_edit['rotation'][0], aoi3D_scene_edit['rotation'][1], pointer_delta_x])

                                # Apply transformation
                                aoi3D_scene_edited = aoi3D_scene.transform(aoi3D_scene_edit['translation'], aoi3D_scene_edit['rotation'])

                                # Write rotation matrix
                                R, _ = cv.Rodrigues(aoi3D_scene_edit['rotation'])
                                cv.putText(video_frame.matrix, f'Rotation matrix:', (20, 160), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv.LINE_AA)
                                cv.putText(video_frame.matrix, f'{R[0][0]:.3f}   {R[0][1]:.3f}   {R[0][2]:.3f}', (40, 200), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv.LINE_AA)
                                cv.putText(video_frame.matrix, f'{R[1][0]:.3f}   {R[1][1]:.3f}   {R[1][2]:.3f}', (40, 240), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv.LINE_AA)
                                cv.putText(video_frame.matrix, f'{R[2][0]:.3f}   {R[2][1]:.3f}   {R[2][2]:.3f}', (40, 280), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv.LINE_AA)
                                
                                # Write translation vector
                                T = aoi3D_scene_edit['translation']
                                cv.putText(video_frame.matrix, f'Translation vector:', (20, 320), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv.LINE_AA)
                                cv.putText(video_frame.matrix, f'{T[0]:.3f}', (40, 360), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv.LINE_AA)
                                cv.putText(video_frame.matrix, f'{T[1]:.3f}', (40, 400), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv.LINE_AA)
                                cv.putText(video_frame.matrix, f'{T[2]:.3f}', (40, 440), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv.LINE_AA)
                            
                            # 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_edited.project(aruco_tracker.get_marker_translation(selected_marker_index), aruco_tracker.get_marker_rotation(selected_marker_index), aruco_camera.get_K())

                            # Draw scene
                            aoi2D_video_scene.draw(video_frame.matrix, pointer, 2, exclude=['Visualisation_Plan'])

                        else:

                            # Write error
                            cv.putText(video_frame.matrix, f'Marker {selected_marker_id} have no AOI scene', (20, 120), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv.LINE_AA)

                    except ValueError:
                        
                        # Write error
                        if selected_marker_id >= 0:
                            cv.putText(video_frame.matrix, f'Marker {selected_marker_id} not found', (20, 120), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv.LINE_AA)

                # Reset left_click
                left_click = (0, 0)

                if args.window:

                    key_pressed = cv.waitKey(1)

                    #if key_pressed != -1:
                    #    print(key_pressed) 

                    # Select previous frame with left arrow
                    if key_pressed == 2:
                        frame_index -= 1

                    # Select next frame with right arrow
                    if key_pressed == 3:
                        frame_index += 1

                    # Clip frame index
                    if frame_index < 0:
                        frame_index = 0

                    # Select coordinate to edit
                    if key_pressed == 49 or key_pressed == 50 or key_pressed == 51:
                        edit_coord = key_pressed - 49

                    # Save selected marker edition using 'Ctrl + s'
                    if key_pressed == 19:

                        if selected_marker_id > 0 and aoi3D_scene_edit != None:

                            aoi_scene_filepath = args.marker_id_scene[f'{selected_marker_id}']
                            aoi3D_scene_edited.save(aoi_scene_filepath)

                            print(f'Saving scene related to marker #{selected_marker_id} into {aoi_scene_filepath}')

                    # Close window using 'Esc' key
                    if key_pressed == 27:
                        break

                    # Reload tracker configuratio on 't' key
                    if key_pressed == 116:
                        load_configuration_file()

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

            # Wait 1 second
            time.sleep(1)

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

        # Stop frame display
        cv.destroyAllWindows()

if __name__ == '__main__':

    main()