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

#!/usr/bin/env python

import argparse
import time

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=None, 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

    # 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_detection = False

    # Init marker selection
    selected_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)

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

            video_capture.set(cv2.CAP_PROP_POS_FRAMES, next_frame_index)

            success, video_frame = video_capture.read()

            if success:

                current_frame_index = video_capture.get(cv2.CAP_PROP_POS_FRAMES) - 1
                current_frame_time = video_capture.get(cv2.CAP_PROP_POS_MSEC)
                
                # Detect markers
                ar_environment.aruco_detector.detect_markers(video_frame)

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

                # Draw focus area
                cv2.rectangle(video_frame, (int(frame_width/6), 0), (int(frame_width*(1-1/6)), int(frame_height)), (255, 150, 150), 1)
                       
                # Write timing
                cv2.rectangle(video_frame, (0, 0), (700, 50), (63, 63, 63), -1)
                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()

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

        # Write detected marker ids
        cv2.putText(video_frame, f'Detected markers: {list(ar_environment.aruco_detector.detected_markers.keys())}', (20, frame_height - 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
                
        # Handle marker selection on left click
        if len(ar_environment.aruco_detector.detected_markers) > 0:

            # Update selected marker id by left clicking on marker
            for (marker_id, marker) in ar_environment.aruco_detector.detected_markers.items():

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

                if marker_aoi.contains_point(left_click):

                    selected_marker_id = marker_id

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

                # Estimate selected marker pose
                ar_environment.aruco_detector.estimate_markers_pose([selected_marker_id])

                # 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)
                           
                    # Apply transformations
                    R = selected_place.rotation.dot(ArUcoScene.make_rotation_matrix(*place_edit[selected_marker_id]['rotation']).T)
                    T = selected_place.translation + numpy.array(place_edit[selected_marker_id]['translation'])

                    edited_place = ArUcoScene.Place(T, R, selected_marker)

                else:

                    edited_place = selected_place
                
                cv2.rectangle(video_frame, (0, 130), (460, 450), (127, 127, 127), -1)
                
                # Write edited rotation matrix
                R = edited_place.rotation
                cv2.putText(video_frame, f'Rotation matrix:', (20, 160), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
                cv2.putText(video_frame, f'{R[0][0]:.3f}   {R[0][1]:.3f}   {R[0][2]:.3f}', (40, 200), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
                cv2.putText(video_frame, f'{R[1][0]:.3f}   {R[1][1]:.3f}   {R[1][2]:.3f}', (40, 240), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv2.LINE_AA)
                cv2.putText(video_frame, f'{R[2][0]:.3f}   {R[2][1]:.3f}   {R[2][2]:.3f}', (40, 280), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv2.LINE_AA)
                
                # Write edited translation vector
                T = edited_place.translation
                cv2.putText(video_frame, f'Translation vector:', (20, 320), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
                cv2.putText(video_frame, f'{T[0]:.3f}', (40, 360), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
                cv2.putText(video_frame, f'{T[1]:.3f}', (40, 400), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv2.LINE_AA)
                cv2.putText(video_frame, f'{T[2]:.3f}', (40, 440), 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

                # Estimate scene pose considering only selected marker
                tvec, rmat, _ = ar_scene.estimate_pose({selected_marker_id: selected_marker})

                # Draw expected marker places
                ar_scene.draw_places(video_frame)

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

                # Draw AOI scene projection with gaze
                aoi_scene_projection.draw_circlecast(video_frame, gaze_position)

            # Catch missing selected marker
            except KeyError:
                
                # Write error
                if selected_marker_id >= 0:

                    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)

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

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

        # Draw pointer
        gaze_position.draw(video_frame)

        # Write selected marker id
        if selected_marker_id >= 0:

            cv2.rectangle(video_frame, (0, 50), (550, 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)

        # Write documentation
        else:
            cv2.rectangle(video_frame, (0, 50), (650, 250), (127, 127, 127), -1)
            cv2.putText(video_frame, f'> Left click on marker: select scene', (20, 80), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)
            cv2.putText(video_frame, f'> T: translate, R: rotate', (20, 120), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)
            cv2.putText(video_frame, f'> Z: switch Z axis edition', (20, 160), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)
            cv2.putText(video_frame, f'> Right click and drag: edit XY axis', (20, 200), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)
            cv2.putText(video_frame, f'> Ctrl + S: save scene', (20, 240), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)

        # Reset left_click
        left_click = (0, 0)

        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

        # 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 detector configuration on 'c' key
        if key_pressed == 99:
            print(f'TODO: Reload ArUcoDetector parameters')
            refresh_detection = True

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

    # Close movie capture
    video_capture.release()

    # Stop frame display
    cv2.destroyAllWindows()

if __name__ == '__main__':

    main()