Working around arscene edition.

1 files changed, 365 insertions, 0 deletions

diff --git a/src/argaze/utils/tobii_segment_arscene_edit.py b/src/argaze/utils/tobii_segment_arscene_edit.py
new file mode 100644
index 0000000..587ca14
--- /dev/null
+++ b/src/argaze/utils/tobii_segment_arscene_edit.py
@@ -0,0 +1,365 @@
+#!/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 tracking 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)
+
+                    # Track markers with pose estimation
+                    ar_scene.aruco_tracker.track(video_frame.matrix)
+
+                else:
+
+                    video_frame = last_frame.copy()
+
+                # Edit fake gaze position from pointer
+                gaze_position = GazeFeatures.GazePosition(pointer, accuracy=2)
+
+                # Copy video frame to edit visualisation on it with out disrupting aruco tracking
+                visu_frame = video_frame.copy()
+
+                # Project scene into frame
+                scene_projection = ar_scene.project(video_frame.matrix, consistent_markers_number=1, visual_hfov=TobiiSpecifications.VISUAL_HFOV, pre_tracked_markers=True)
+
+                # Draw tracked markers
+                ar_scene.aruco_tracker.draw_tracked_markers(visu_frame.matrix)
+
+                # Draw scene projection
+                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_tracker.tracked_markers) > 0:
+
+                    # Write detected marker ids
+                    cv.putText(visu_frame.matrix, f'Detected markers: {list(ar_scene.aruco_tracker.tracked_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_tracker.tracked_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_tracker.tracked_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 tracker 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()
\ No newline at end of file