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

#!/usr/bin/env python

""" """

__author__ = "Théo de la Hogue"
__credits__ = []
__copyright__ = "Copyright 2023, Ecole Nationale de l'Aviation Civile (ENAC)"
__license__ = "BSD"

import argparse
import time
import itertools

from argaze import ArFeatures, GazeFeatures
from argaze.AreaOfInterest import AOIFeatures
from argaze.ArUcoMarkers import ArUcoScene
from argaze.utils import MiscFeatures

from tobiiproglasses2 import *

import cv2
import numpy

def main():
    """
    Load AR environment from .json file, detect ArUco markers into movie frames and estimate environment pose.
    Edit environment setup to improve pose estimation.
    """

    # Manage arguments
    parser = argparse.ArgumentParser(description=main.__doc__.split('-')[0])
    parser.add_argument('environment', metavar='ENVIRONMENT', type=str, help='ar environment filepath')
    parser.add_argument('movie', metavar='MOVIE', type=str, default=None, help='movie path')
    parser.add_argument('-s','--start', metavar='START', type=float, default=0., help='start time in second')
    parser.add_argument('-o', '--output', metavar='OUT', type=str, default='environment.json', help='edited ar environment file path')
    args = parser.parse_args()

    # Load AR enviroment
    ar_environment = ArFeatures.ArEnvironment.from_json(args.environment)

    #print('ArEnvironment:\n', ar_environment)

    # Select first AR scene
    ar_scene = list(ar_environment.scenes.values())[0]

    # Create a window to display AR environment
    cv2.namedWindow(ar_environment.name, cv2.WINDOW_AUTOSIZE)

    # Init mouse interaction
    pointer = (0, 0)
    left_click = (0, 0)
    right_click = (0, 0)
    right_drag = (0, 0)
    right_button = False
    edit_trans = False # translate
    edit_z = False
    snap = False
    draw_help = False
    draw_grid = False
    draw_cover = False
    pose_mode = 0
    z_grid = 100.

    # Update pointer position
    def on_mouse_event(event, x, y, flags, param):

        nonlocal pointer
        nonlocal left_click
        nonlocal right_click
        nonlocal right_drag
        nonlocal right_button

        # Update pointer
        pointer = (x, y)

        # Update left_click
        if event == cv2.EVENT_LBUTTONUP:

            left_click = pointer

        # Udpate right_button
        elif event == cv2.EVENT_RBUTTONDOWN and not right_button:

            right_button = True
            right_click = pointer

        elif event == cv2.EVENT_RBUTTONUP and right_button:

            right_button = False

        # Udpate right_drag
        if right_button:

            right_drag = (pointer[0] - right_click[0], pointer[1] - right_click[1])

    # Attach mouse callback to window
    cv2.setMouseCallback(ar_environment.name, on_mouse_event)

    # Enable movie video capture
    video_capture = cv2.VideoCapture(args.movie)

    video_fps = video_capture.get(cv2.CAP_PROP_FPS)
    frame_width = int(video_capture.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT))

    # Enable exit signal handler
    exit = MiscFeatures.ExitSignalHandler()

    # Init frame selection
    current_frame_index = -1
    _, current_frame = video_capture.read()
    next_frame_index = int(args.start * video_fps)
    refresh = False

    # Init marker selection
    selected_marker_id = -1
    hovered_marker_id = -1

    # Init place edition
    place_edit = {}
     
    while not exit.status():

        # Edit fake gaze position from pointer
        gaze_position = GazeFeatures.GazePosition(pointer, precision=2)

        # Reset info frame
        info_frame = numpy.full((850, 1500, 3), 0, dtype=numpy.uint8)

        # Select a new frame and detect markers once
        if next_frame_index != current_frame_index or refresh or draw_cover:

            video_capture.set(cv2.CAP_PROP_POS_FRAMES, next_frame_index)

            success, video_frame = video_capture.read()

            if success:

                # Refresh once
                refresh = False

                current_frame_index = video_capture.get(cv2.CAP_PROP_POS_FRAMES) - 1
                current_frame_time = video_capture.get(cv2.CAP_PROP_POS_MSEC)

                # Hide zone
                if draw_cover:

                    # Draw black circle under pointer
                    cv2.circle(video_frame, pointer, 50, (0, 0, 0), -1)
                
                # Detect markers
                ar_environment.aruco_detector.detect_markers(video_frame)

                # Estimate all marker's pose
                ar_environment.aruco_detector.estimate_markers_pose()

                # Write detected markers
                cv2.putText(video_frame, f'{list(ar_environment.aruco_detector.detected_markers.keys())}', (20, frame_height-80), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
                
                # Draw gray panel on top
                cv2.rectangle(video_frame, (0, 0), (frame_width, 50), (63, 63, 63), -1)

                # Draw camera calibration
                if draw_grid:

                    cv2.putText(video_frame, f'Grid at {z_grid} cm', (500, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
                    ar_environment.aruco_detector.optic_parameters.draw(video_frame, frame_width/10, frame_height/10, z_grid, color=(127, 127, 127))

                # Write timing
                cv2.putText(video_frame, f'Time: {int(current_frame_time)} ms', (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
                
                # Copy frame
                current_frame = video_frame.copy()

        # Keep last frame
        else:

            video_frame = current_frame.copy()

        # Handle marker selection on left click
        if len(ar_environment.aruco_detector.detected_markers) > 0:

            # Update selected and hovered markers id
            selected_marker_id = -1
            hovered_marker_id = -1
            for (marker_id, marker) in ar_environment.aruco_detector.detected_markers.items():

                marker_aoi = marker.corners.reshape(4, 2).view(AOIFeatures.AreaOfInterest)

                # Select by left clicking on marker
                if marker_aoi.contains_point(left_click):

                    selected_marker_id = marker_id

                # Hover by pointing on marker
                if marker_aoi.contains_point(pointer):

                    hovered_marker_id = marker_id

            # Edit marker's color
            color_list = list(itertools.permutations([0, 255, 255]))

            for i, m in ar_environment.aruco_detector.detected_markers.items():

                m.color = color_list[i%len(color_list)]

                if i == selected_marker_id or i == hovered_marker_id:
                    continue

                if hovered_marker_id > 0:
                    m.color = (0, 0, 0)
                else:
                    m.color = (127, 127, 127)

                # Draw center
                cv2.circle(video_frame, m.center.astype(int), 5, m.color, -1)

            try:

                # A marker is selected
                if selected_marker_id >= 0:

                    try:
                        
                        # Retreive selected marker
                        selected_marker = ar_environment.aruco_detector.detected_markers[selected_marker_id]

                        # Write selected marker id
                        cv2.rectangle(info_frame, (0, 0), (500, 50), selected_marker.color, -1)
                        cv2.putText(info_frame, f'Selected marker #{selected_marker.identifier}', (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 1, cv2.LINE_AA)
                        cv2.rectangle(info_frame, (0, 50), (500, frame_height), (255, 255, 255), -1)

                        # Write selected marker rotation matrix
                        R = ArUcoScene.make_euler_rotation_vector(selected_marker.rotation)
                        cv2.putText(info_frame, f'Rotation (camera axis)', (20, 120), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 1, cv2.LINE_AA)
                        cv2.putText(info_frame, f'{R[0]:.3f}', (40, 160), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
                        cv2.putText(info_frame, f'{R[1]:.3f}', (40, 200), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv2.LINE_AA)
                        cv2.putText(info_frame, f'{R[2]:.3f}', (40, 240), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv2.LINE_AA)
                        
                        # Write selected marker translation vector
                        T = selected_marker.translation
                        cv2.putText(info_frame, f'Translation (camera axis):', (20, 320), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 1, cv2.LINE_AA)
                        cv2.putText(info_frame, f'{T[0]:.3f}', (40, 360), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
                        cv2.putText(info_frame, f'{T[1]:.3f}', (40, 400), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv2.LINE_AA)
                        cv2.putText(info_frame, f'{T[2]:.3f}', (40, 440), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv2.LINE_AA)

                        # Retreive selected marker place
                        selected_place = ar_scene.aruco_scene.places[selected_marker_id]

                        # On right click
                        if right_button:

                            pointer_delta_x, pointer_delta_y = right_drag[0] / frame_width, right_drag[1] / frame_height

                            place_edit[selected_marker_id] = {'rotation': (0, 0, 0), 'translation': (0, 0, 0)}

                            if edit_trans:
                            
                                # Edit place rotation
                                if edit_z:
                                    place_edit[selected_marker_id]['rotation'] = (0, 0, -pointer_delta_y)
                                else:
                                    place_edit[selected_marker_id]['rotation'] = (pointer_delta_y, pointer_delta_x, 0)

                            else:

                                # Edit place translation
                                if edit_z:
                                     place_edit[selected_marker_id]['translation'] = (0, 0, pointer_delta_y)
                                else:
                                    place_edit[selected_marker_id]['translation'] = (-pointer_delta_x, pointer_delta_y, 0)
                                   
                            # Edit transformations
                            R = ArUcoScene.make_rotation_matrix(*place_edit[selected_marker_id]['rotation']).T
                            T = numpy.array(place_edit[selected_marker_id]['translation'])

                            # Apply transformations
                            edited_place = ArUcoScene.Place(selected_place.translation + T, selected_place.rotation.dot(R), selected_marker)

                        else:

                            edited_place = selected_place

                        # A marker is hovered while another is selected
                        if hovered_marker_id >= 0 and hovered_marker_id != selected_marker_id:

                            # Retreive hovered marker
                            hovered_marker = ar_environment.aruco_detector.detected_markers[hovered_marker_id]

                            # Write hovered marker id
                            cv2.rectangle(info_frame, (500, 0), (1000, 50), hovered_marker.color, -1)
                            cv2.putText(info_frame, f'Hovered marker #{hovered_marker.identifier}', (520, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 1, cv2.LINE_AA)
                            cv2.rectangle(info_frame, (500, 50), (1000, frame_height), (255, 255, 255), -1)

                            # Write hovered marker rotation matrix
                            R = ArUcoScene.make_euler_rotation_vector(hovered_marker.rotation)
                            cv2.putText(info_frame, f'Rotation (camera axis)', (520, 120), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{R[0]:.3f}', (540, 160), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{R[1]:.3f}', (540, 200), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{R[2]:.3f}', (540, 240), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv2.LINE_AA)
                            
                            # Write hovered marker translation vector
                            T = hovered_marker.translation
                            cv2.putText(info_frame, f'Translation (camera axis):', (520, 320), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{T[0]:.3f}', (540, 360), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{T[1]:.3f}', (540, 400), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{T[2]:.3f}', (540, 440), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv2.LINE_AA)

                            # Retreive hovered marker place
                            hovered_place = ar_scene.aruco_scene.places[hovered_marker_id]

                            # Write hovered place rotation matrix
                            R = ArUcoScene.make_euler_rotation_vector(hovered_place.rotation)
                            cv2.putText(info_frame, f'Rotation (scene axis):', (520, 500), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{R[0]:.3f}', (540, 540), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{R[1]:.3f}', (540, 580), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{R[2]:.3f}', (540, 620), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv2.LINE_AA)
                            
                            # Write hovered place translation vector
                            T = hovered_place.translation
                            cv2.putText(info_frame, f'Translation (scene axis):', (520, 700), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{T[0]:.3f}', (540, 740), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{T[1]:.3f}', (540, 780), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{T[2]:.3f}', (540, 820), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv2.LINE_AA)

                            # Rotation between markers and places
                            markers_rotation_matrix = hovered_marker.rotation.dot(selected_marker.rotation.T)
                            places_rotation_matrix = hovered_place.rotation.dot(selected_place.rotation.T)

                            markers_rotation_vector = ArUcoScene.make_euler_rotation_vector(markers_rotation_matrix)
                            places_rotation_vector = ArUcoScene.make_euler_rotation_vector(places_rotation_matrix)

                            # Translation info between markers and places
                            markers_translation = hovered_marker.translation - selected_marker.translation
                            places_translation = hovered_place.translation - selected_place.translation

                            # Write selected/hovered markers id
                            cv2.rectangle(info_frame, (1000, 0), (1500, 50), (63, 63, 63), -1)
                            cv2.putText(info_frame, f'#{selected_marker.identifier} -> #{hovered_marker.identifier}', (1020, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
                            cv2.rectangle(info_frame, (1000, 50), (1500, frame_height), (190, 190, 190), -1)

                            # Write selected/hovered markers rotation matrix
                            R = markers_rotation_vector
                            cv2.putText(info_frame, f'Rotation (camera axis)', (1020, 120), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{R[0]:.3f}', (1040, 160), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{R[1]:.3f}', (1040, 200), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{R[2]:.3f}', (1040, 240), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv2.LINE_AA)
                            
                            # Write selected/hovered markers translation vector
                            T = markers_translation
                            cv2.putText(info_frame, f'Translation (camera axis):', (1020, 320), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{T[0]:.3f}', (1040, 360), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{T[1]:.3f}', (1040, 400), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{T[2]:.3f}', (1040, 440), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv2.LINE_AA)

                            # Write selected/hovered places rotation matrix
                            R = places_rotation_vector
                            cv2.putText(info_frame, f'Rotation (scene axis):', (1020, 500), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{R[0]:.3f}', (1040, 540), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{R[1]:.3f}', (1040, 580), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{R[2]:.3f}', (1040, 620), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv2.LINE_AA)
                            
                            # Write selected/hovered places translation vector
                            T = places_translation
                            cv2.putText(info_frame, f'Translation (scene axis):', (1020, 700), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{T[0]:.3f}', (1040, 740), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{T[1]:.3f}', (1040, 780), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv2.LINE_AA)
                            cv2.putText(info_frame, f'{T[2]:.3f}', (1040, 820), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv2.LINE_AA)

                            if snap: 

                                # Snap once
                                snap = False

                                print(f'******* SNAP {selected_marker_id} / {hovered_marker_id} *******')

                                '''
                                # Edit rotation transformation
                                #R = places_rotation_matrix.dot(rmat.T).dot(markers_rotation_matrix.T).dot(rmat)

                                #rmat_places_rotation_vector = ArUcoScene.make_euler_rotation_vector(places_rotation_matrix.dot(rmat.T))
                                rdiff = places_rotation_vector - markers_rotation_vector
                                R = ArUcoScene.make_rotation_matrix(*rdiff)

                                print(f'markers_rotation_vector: {markers_rotation_vector}')
                                print(f'places_rotation_vector: {places_rotation_vector}')
                                print(f'rdiff: {rdiff}')
                                print(f'R: {ArUcoScene.make_euler_rotation_vector(R)}')
                                '''

                                # Edit translation transformation
                                T = (places_translation.dot(rmat.T) - markers_translation).dot(rmat)

                                print(f'markers_translation: {markers_translation} ({numpy.linalg.norm(markers_translation)})')
                                print(f'places_translation: {places_translation} ({numpy.linalg.norm(places_translation)})')
                                print(f'T: {T} ({numpy.linalg.norm(T)})')

                                # Apply transformations
                                edited_place = ArUcoScene.Place(selected_place.translation + T, selected_place.rotation, selected_marker)
                        
                        # Write edited place rotation matrix
                        R = ArUcoScene.make_euler_rotation_vector(edited_place.rotation)
                        cv2.putText(info_frame, f'Rotation (scene axis):', (20, 500), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 1, cv2.LINE_AA)
                        cv2.putText(info_frame, f'{R[0]:.3f}', (40, 540), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
                        cv2.putText(info_frame, f'{R[1]:.3f}', (40, 580), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv2.LINE_AA)
                        cv2.putText(info_frame, f'{R[2]:.3f}', (40, 620), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv2.LINE_AA)
                        
                        # Write edited place translation vector
                        T = edited_place.translation
                        cv2.putText(info_frame, f'Translation (scene axis):', (20, 700), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 1, cv2.LINE_AA)
                        cv2.putText(info_frame, f'{T[0]:.3f}', (40, 740), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
                        cv2.putText(info_frame, f'{T[1]:.3f}', (40, 780), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv2.LINE_AA)
                        cv2.putText(info_frame, f'{T[2]:.3f}', (40, 820), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv2.LINE_AA)

                        # Replace selected place by edited place
                        ar_scene.aruco_scene.places[selected_marker_id] = edited_place

                        # Refresh places consistency
                        ar_scene.aruco_scene.init_places_consistency()

                        # Estimate scene pose from each marker
                        cv2.putText(video_frame, f'Single marker scene pose estimation', (20, frame_height - 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)

                        for i, m in ar_environment.aruco_detector.detected_markers.items():

                            tvec, rmat = ar_scene.aruco_scene.estimate_pose_from_single_marker(m)

                            # Project AOI scene into frame according estimated pose
                            aoi_scene_projection = ar_scene.project(tvec, rmat, visual_hfov=TobiiSpecifications.VISUAL_HFOV)

                            if i == selected_marker_id:

                                # Draw AOI scene projection with gaze
                                aoi_scene_projection.draw_circlecast(video_frame, gaze_position, base_color=m.color, looked_color=(255, 255, 255))

                            else:

                                # Draw AOI scene
                                aoi_scene_projection.draw(video_frame, color=m.color)

                        # Draw expected marker places
                        ar_scene.draw_places(video_frame)

                    # Catch missing selected marker
                    except KeyError:
                        
                        cv2.putText(video_frame, f'Marker {selected_marker_id} not found', (20, 120), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)

                # No marker selected
                else:

                    cv2.putText(info_frame, f'Left click on marker to select it', (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)
                    
                    if len(ar_environment.aruco_detector.detected_markers) > 1:

                        # Check markers consistency
                        consistent_markers, unconsistent_markers, unconsistencies = ar_scene.aruco_scene.check_markers_consistency(ar_environment.aruco_detector.detected_markers, ar_scene.angle_tolerance, ar_scene.distance_tolerance)

                        # No marker hovered
                        if hovered_marker_id < 0:

                            # Set unconsistent marker color to red 
                            for i, m in ar_environment.aruco_detector.detected_markers.items():
                                if i in list(unconsistent_markers.keys()) and i != hovered_marker_id:
                                    m.color = (0, 0, 255)

                        # Write unconsistencies
                        line = 0
                        for i, (label, value) in enumerate(unconsistencies['rotation'].items()):

                            current_rotation = value['current']
                            expected_rotation = value['expected']

                            cv2.putText(info_frame, f'Unconsistent rotation {label}: [{current_rotation[0]:.3f} {current_rotation[1]:.3f} {current_rotation[2]:.3f}]', (20, 120+line*40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
                            line += 1

                            cv2.putText(info_frame, f'Expected rotation {label}: [{expected_rotation[0]:.3f} {expected_rotation[1]:.3f} {expected_rotation[2]:.3f}]', (20, 120+line*40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
                            line += 1

                        for i, (label, value) in enumerate(unconsistencies['translation'].items()):

                            current_translation = value['current']
                            expected_translation = value['expected']

                            cv2.putText(info_frame, f'Unconsistent translation {label}: {current_translation:.3f}', (20, 120+ line*40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
                            line += 1

                            cv2.putText(info_frame, f'Expected translation {label}: {expected_translation:.3f}', (20, 120+ line*40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
                            line += 1

                    # Force pose mode to single marker scene pose estimation
                    else: 

                        pose_mode = 0
                    
                    # Single marker scene pose estimation
                    if pose_mode == 0:

                        cv2.putText(video_frame, f'Single marker scene pose estimation', (20, frame_height - 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
                
                        for i, m in ar_environment.aruco_detector.detected_markers.items():

                            tvec, rmat = ar_scene.aruco_scene.estimate_pose_from_single_marker(m)

                            # Project AOI scene into frame according estimated pose
                            aoi_scene_projection = ar_scene.project(tvec, rmat, visual_hfov=TobiiSpecifications.VISUAL_HFOV)

                            # Draw AOI scene
                            aoi_scene_projection.draw(video_frame, color=m.color)

                    # Consistent markers scene pose estimation
                    if pose_mode == 1:

                        cv2.putText(video_frame, f'Consistent markers scene pose estimation', (20, frame_height - 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
                
                        tvec, rmat = ar_scene.aruco_scene.estimate_pose_from_markers(consistent_markers)

                        # Project AOI scene into frame according estimated pose
                        aoi_scene_projection = ar_scene.project(tvec, rmat, visual_hfov=TobiiSpecifications.VISUAL_HFOV)

                        # Draw AOI scene
                        aoi_scene_projection.draw(video_frame, color=(255, 255, 255))

                    # ArUco marker axis scene pose estimation
                    elif pose_mode == 2:

                        # Write pose estimation strategy
                        cv2.putText(video_frame, f'ArUco marker axis scene pose estimation', (20, frame_height - 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
                
                        for axis_name, axis_markers in ar_scene.aruco_axis.items():

                            try:

                                origin_marker = ar_environment.aruco_detector.detected_markers[axis_markers['origin_marker']]
                                horizontal_axis_marker = ar_environment.aruco_detector.detected_markers[axis_markers['horizontal_axis_marker']]
                                vertical_axis_marker = ar_environment.aruco_detector.detected_markers[axis_markers['vertical_axis_marker']]

                                tvec, rmat = ar_scene.aruco_scene.estimate_pose_from_axis_markers(origin_marker, horizontal_axis_marker, vertical_axis_marker)

                                # Project AOI scene into frame according estimated pose
                                aoi_scene_projection = ar_scene.project(tvec, rmat, visual_hfov=TobiiSpecifications.VISUAL_HFOV)

                                # Draw AOI scene
                                aoi_scene_projection.draw(video_frame, color=(255, 255, 255))

                                break

                            except:
                                pass

                    # ArUco AOI scene building
                    elif pose_mode == 3:

                        # Write pose estimation strategy
                        cv2.putText(video_frame, f'ArUco AOI scene building', (20, frame_height - 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
                       
                        try :

                            # Try to build AOI scene from detected ArUco marker corners
                            aoi_scene_projection = ar_scene.build_aruco_aoi_scene(ar_environment.aruco_detector.detected_markers)

                            # Draw AOI scene
                            aoi_scene_projection.draw(video_frame, color=(255, 255, 255))

                        except:
                            pass

                        # Draw expected marker places
                        #ar_scene.draw_places(video_frame)

            # Catch exceptions raised by estimate_pose and project methods
            except (ArFeatures.PoseEstimationFailed, ArFeatures.SceneProjectionFailed) as e:

                cv2.rectangle(video_frame, (0, 90), (700, 130), (127, 127, 127), -1)
                cv2.putText(video_frame, f'Error: {e}', (20, 100), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)
            
                # Draw frame
                cv2.imshow(ar_environment.name, video_frame)

        # Draw detected markers
        ar_environment.aruco_detector.draw_detected_markers(video_frame)

        # Draw pointer
        gaze_position.draw(video_frame)

        # Write documentation
        cv2.putText(video_frame, f'Press \'h\' for help', (950, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)

        if draw_help:

            cv2.rectangle(video_frame, (0, 50), (700, 300), (127, 127, 127), -1)
            cv2.putText(video_frame, f'> Left click on marker: select marker', (20, 80), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)
            cv2.putText(video_frame, f'> Left click on frame: unselect marker', (20, 120), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)
            cv2.putText(video_frame, f'> T: translate, R: rotate, Z: select axis', (20, 160), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)
            cv2.putText(video_frame, f'> Right click and drag: edit axis', (20, 200), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)
            cv2.putText(video_frame, f'> Ctrl + S: save environment', (20, 240), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)
            cv2.putText(video_frame, f'> Backspace: reload environment', (20, 280), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)

        # Write selected marker id
        elif selected_marker_id >= 0:

            cv2.rectangle(video_frame, (0, 50), (700, 90), (127, 127, 127), -1)

            # Select color
            if edit_z:
                str_axis = 'Z'
                color_axis = (255, 0, 0)
            else:
                str_axis = 'XY'
                color_axis = (0, 255, 255)

            if edit_trans:
                cv2.putText(video_frame, f'Rotate marker {selected_marker_id} around axis {str_axis}', (20, 80), cv2.FONT_HERSHEY_SIMPLEX, 1, color_axis, 1, cv2.LINE_AA)
            else:
                cv2.putText(video_frame, f'Translate marker {selected_marker_id} along axis {str_axis}', (20, 80), cv2.FONT_HERSHEY_SIMPLEX, 1, color_axis, 1, cv2.LINE_AA)

        key_pressed = cv2.waitKey(10)

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

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

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

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

        # Edit rotation with r key
        if key_pressed == 114:
            edit_trans = True

        # Edit translation with t key
        if key_pressed == 116:
            edit_trans = False

        # Switch Z axis edition
        if key_pressed == 122:
            edit_z = not edit_z

        # Snap hovered marker with s key
        if key_pressed == 115:
            snap = True

        # Switch help mode with h key
        if key_pressed == 104:
            draw_help = not draw_help

        # Switch grid mode with g key
        if key_pressed == 103:
            draw_grid = not draw_grid
            refresh = True

        # Raise z grid with down arrow
        if key_pressed == 0:
            z_grid += 10.
            refresh = True

        # Unraise z grid with up arrow
        if key_pressed == 1:
            z_grid -= 10.
            refresh = True

        # Switch draw_cover mode with c key
        if key_pressed == 99:
            draw_cover = not draw_cover

        # Switch pose estimation mode with m key
        if key_pressed == 109:
            pose_mode += 1
            if pose_mode > 3:
                pose_mode = 0

        # Save selected marker edition using 'Ctrl + s'
        if key_pressed == 19:
            ar_environment.to_json(args.output)
            print(f'Environment saved into {args.output}')

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

        # Reload environment on 'Backspace' key
        if key_pressed == 127:
            ar_environment = ArFeatures.ArEnvironment.from_json(args.environment)
            print(f'Environment reloaded from {args.environment}')
            refresh = True

        # Display video
        cv2.imshow(ar_environment.name, video_frame)

        # Display info
        cv2.imshow('Info', info_frame)

    # Close movie capture
    video_capture.release()

    # Stop frame display
    cv2.destroyAllWindows()

if __name__ == '__main__':

    main()