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

#!/usr/bin/env python

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
    cover = False
    draw_help = False
    pose_mode = 0

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

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

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

            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)

                # Hide zone
                if 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)

                # 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), (frame_width, 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()

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

                else:

                    selected_marker_id = -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.putText(info_frame, f'Selected marker #{selected_marker.identifier}', (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)

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

                        # Write selected marker rotation matrix
                        R = selected_marker.rotation
                        cv2.putText(info_frame, f'Rotation matrix:', (20, 120), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
                        cv2.putText(info_frame, f'{R[0][0]:.3f}   {R[0][1]:.3f}   {R[0][2]:.3f}', (40, 160), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
                        cv2.putText(info_frame, f'{R[1][0]:.3f}   {R[1][1]:.3f}   {R[1][2]:.3f}', (40, 200), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv2.LINE_AA)
                        cv2.putText(info_frame, f'{R[2][0]:.3f}   {R[2][1]:.3f}   {R[2][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 vector:', (20, 320), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 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]

                        cv2.putText(info_frame, f'Edited place #{selected_place.marker.identifier}', (520, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)

                        # 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
                        
                        # Write edited place rotation matrix
                        R = edited_place.rotation
                        cv2.putText(info_frame, f'Rotation matrix:', (520, 120), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
                        cv2.putText(info_frame, f'{R[0][0]:.3f}   {R[0][1]:.3f}   {R[0][2]:.3f}', (540, 160), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
                        cv2.putText(info_frame, f'{R[1][0]:.3f}   {R[1][1]:.3f}   {R[1][2]:.3f}', (540, 200), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv2.LINE_AA)
                        cv2.putText(info_frame, f'{R[2][0]:.3f}   {R[2][1]:.3f}   {R[2][2]:.3f}', (540, 240), 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 vector:', (520, 320), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 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)

                        # 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, strategy, _ = ar_scene.estimate_pose({selected_marker_id: selected_marker})
                        strategy = strategy.replace('_', ' ')

                        # Write pose estimation strategy
                        cv2.putText(video_frame, f'{strategy} with marker {selected_marker_id}', (20, frame_height - 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
            
                        # 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:
                        
                        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, (255, 255, 255), 1, cv2.LINE_AA)

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

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

                    # Edit marker's color
                    for i, m in ar_environment.aruco_detector.detected_markers.items():

                        m.color = list(itertools.permutations([0, 255, 255]))[i]

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

                        # Set unconsistent marker color to red 
                        if i in list(unconsistent_markers.keys()):
                            m.color = (0, 0, 255)
 
                    # 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)

                            aoi_scene_projection.draw(video_frame, color=m.color)

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

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

                                aoi_scene_projection.draw(video_frame, color=(255, 0, 0))

                                break

                            except:
                                pass

                    # ArUco AOI scene building
                    elif pose_mode == 2:

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

                        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, 260), (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)

        # 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

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

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

        # Switch pose estimation mode with m key
        if key_pressed == 109:
            pose_mode += 1
            if pose_mode > 2:
                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 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)

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

    # Close movie capture
    video_capture.release()

    # Stop frame display
    cv2.destroyAllWindows()

if __name__ == '__main__':

    main()