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

#!/usr/bin/env python

import argparse
import os
import json
import time

from argaze import *
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():
    """
    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('-t', '--time_range', metavar=('START_TIME', 'END_TIME'), nargs=2, type=float, default=(0., None), help='start and end time (in second)')
    parser.add_argument('-p', '--project_path', metavar='ARGAZE_PROJECT', type=str, default=None, help='json argaze project filepath')
    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 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}/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.duration/1e6} s\n\twidth: {tobii_segment_video.width} px\n\theight: {tobii_segment_video.height} px')

        # Load ar scene
        ar_scene = ArScene.ArScene.from_json(args.project_path)

        print(ar_scene)

        # Display first frame
        video_ts, video_frame = tobii_segment_video.get_frame(0)
        cv.imshow(f'Segment {tobii_segment.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_trans = False # translate
        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.id} ArUco marker editor', on_mouse_event)

        # Frame selector loop
        frame_index = 0
        last_frame_index = -1
        last_frame = video_frame.copy()
        force_update = False

        selected_marker_id = -1

        try:

            while True:

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

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

                    last_frame_index = frame_index
                    last_frame = video_frame.copy()

                    # Hide frame left and right borders before detection 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)

                    try:

                        # Estimate scene pose from ArUco markers into frame.
                        tvec, rmat, _ = ar_scene.estimate_pose(video_frame.matrix)

                    # Catch exceptions raised by estimate_pose method
                    except ArScene.PoseEstimationFailed as e:

                        cv.rectangle(visu_frame.matrix, (0, 100), (550, 150), (127, 127, 127), -1)
                        cv.putText(visu_frame.matrix, str(e), (20, 130), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv.LINE_AA)

                else:

                    video_frame = last_frame.copy()

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

                # Copy video frame to edit visualisation on it with out disrupting aruco detection
                visu_frame = video_frame.copy()

                try:

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

                # Catch exceptions raised by project method
                except ArScene.SceneProjectionFailed as e:

                    cv.rectangle(visu_frame.matrix, (0, 100), (550, 150), (127, 127, 127), -1)
                    cv.putText(visu_frame.matrix, str(e), (20, 130), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv.LINE_AA)

                # Draw detected markers
                ar_scene.aruco_detector.draw_detected_markers(visu_frame.matrix)

                # Draw scene projection
                aoi_scene_projection.draw(visu_frame.matrix, (0, 0), color=(0, 255, 255))

                # Project 3D scene on each video frame and the visualisation frame
                if len(ar_scene.aruco_detector.detected_markers) > 0:

                    # Write detected marker ids
                    cv.putText(visu_frame.matrix, f'Detected markers: {list(ar_scene.aruco_detector.detected_markers.keys())}', (20, visu_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 (marker_id, marker) in ar_scene.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 marker index
                        selected_marker = ar_scene.aruco_detector.detected_markers[selected_marker_id]

                        marker_x, marker_y = selected_marker.center
                        '''
                        if right_button:

                            pointer_delta_x, pointer_delta_y = (right_click[0] - marker_x) / (visu_frame.width/3), (marker_y - right_click[1]) / (visu_frame.width/3)

                            if edit_trans:
                            
                                # 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], -1*pointer_delta_y])

                            else:

                                # Edit scene translation
                                if edit_coord == 0:
                                    aoi3D_scene_edit['translation'] = numpy.array([pointer_delta_x, aoi3D_scene_edit['translation'][1], aoi3D_scene_edit['translation'][2]])

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

                                elif edit_coord == 2:
                                    aoi3D_scene_edit['translation'] = numpy.array([aoi3D_scene_edit['translation'][0], aoi3D_scene_edit['translation'][1], 2*pointer_delta_y])
                        '''
                        # Apply transformation
                        aoi_scene_edited = ar_scene.aoi_scene#.transform(aoi3D_scene_edit['translation'], aoi3D_scene_edit['rotation'])
                        
                        cv.rectangle(visu_frame.matrix, (0, 130), (460, 450), (127, 127, 127), -1)
                        '''
                        # Write rotation matrix
                        R, _ = cv.Rodrigues(aoi3D_scene_edit['rotation'])
                        cv.putText(visu_frame.matrix, f'Rotation matrix:', (20, 160), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv.LINE_AA)
                        cv.putText(visu_frame.matrix, f'{R[0][0]:.3f}   {R[0][1]:.3f}   {R[0][2]:.3f}', (40, 200), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv.LINE_AA)
                        cv.putText(visu_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(visu_frame.matrix, f'{R[2][0]:.3f}   {R[2][1]:.3f}   {R[2][2]:.3f}', (40, 280), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv.LINE_AA)
                        
                        # Write translation vector
                        T = aoi3D_scene_edit['translation']
                        cv.putText(visu_frame.matrix, f'Translation vector:', (20, 320), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv.LINE_AA)
                        cv.putText(visu_frame.matrix, f'{T[0]:.3f}', (40, 360), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv.LINE_AA)
                        cv.putText(visu_frame.matrix, f'{T[1]:.3f}', (40, 400), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv.LINE_AA)
                        cv.putText(visu_frame.matrix, f'{T[2]:.3f}', (40, 440), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 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.
                        scene_projection_edited = aoi_scene_edited.project(selected_marker.translation, selected_marker.rotation, ar_scene.aruco_camera.K)

                        # Draw aoi scene
                        scene_projection_edited.draw_raycast(visu_frame.matrix, gaze_position)

                    # Write warning related to marker pose processing 
                    except UserWarning as e:

                        cv.putText(visu_frame.matrix, f'Marker {selected_marker_id}: {e}', (20, 120), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv.LINE_AA)

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

                # Draw focus area
                cv.rectangle(visu_frame.matrix, (int(visu_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)
                 
                # Draw pointer
                gaze_position.draw(visu_frame.matrix)

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

                # Write selected marker id
                if selected_marker_id >= 0:

                    cv.rectangle(visu_frame.matrix, (0, 50), (550, 90), (127, 127, 127), -1)

                    # Select color
                    if edit_coord == 0:
                        color_axis = (0, 0, 255)

                    elif edit_coord == 1:
                        color_axis = (0, 255, 0)
                        
                    elif edit_coord == 2:
                        color_axis = (255, 0, 0)

                    if edit_trans:
                        cv.putText(visu_frame.matrix, f'Rotate marker {selected_marker_id} around axis {edit_coord + 1}', (20, 80), cv.FONT_HERSHEY_SIMPLEX, 1, color_axis, 1, cv.LINE_AA)
                    else:
                        cv.putText(visu_frame.matrix, f'Translate marker {selected_marker_id} along axis {edit_coord + 1}', (20, 80), cv.FONT_HERSHEY_SIMPLEX, 1, color_axis, 1, cv.LINE_AA)

                # Write documentation
                else:
                    cv.rectangle(visu_frame.matrix, (0, 50), (650, 250), (127, 127, 127), -1)
                    cv.putText(visu_frame.matrix, f'> Left click on marker: select scene', (20, 80), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv.LINE_AA)
                    cv.putText(visu_frame.matrix, f'> T: translate, R: rotate', (20, 120), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv.LINE_AA)
                    cv.putText(visu_frame.matrix, f'> Shift + 0/1/2: select axis', (20, 160), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv.LINE_AA)
                    cv.putText(visu_frame.matrix, f'> Right click and drag: edit axis', (20, 200), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv.LINE_AA)
                    cv.putText(visu_frame.matrix, f'> Ctrl + S: save scene', (20, 240), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 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

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

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

                    # Select coordinate to edit with Shift + 0, 1 or 2
                    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}']
                            aoi_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 detecter configuration on 'c' key
                    if key_pressed == 99:
                        load_configuration_file()
                        force_update = True

                    # Display video
                    cv.imshow(f'Segment {tobii_segment.id} ArUco marker editor', visu_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()