Adding a new utils script to edit aruco pose.

1 files changed, 305 insertions, 0 deletions

diff --git a/src/argaze/utils/edit_tobii_segment_aruco_pose.py b/src/argaze/utils/edit_tobii_segment_aruco_pose.py
new file mode 100644
index 0000000..fa0675d
--- /dev/null
+++ b/src/argaze/utils/edit_tobii_segment_aruco_pose.py
@@ -0,0 +1,305 @@
+#!/usr/bin/env python
+
+import argparse
+import os
+import json
+import time
+
+from argaze import DataStructures
+from argaze import GazeFeatures
+from argaze.TobiiGlassesPro2 import TobiiEntities, TobiiVideo
+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('-r', '--time_range', metavar=('START_TIME', 'END_TIME'), nargs=2, type=float, default=(0., None), help='start and end time (in second)')
+    parser.add_argument('-c', '--camera_calibration', metavar='CAM_CALIB', type=str, default=None, help='json camera calibration filepath')
+    parser.add_argument('-p', '--aruco_tracker_configuration', metavar='TRACK_CONFIG', type=str, default=None, help='json aruco tracker configuration filepath')
+    parser.add_argument('-md', '--marker_dictionary', metavar='MARKER_DICT', type=str, default='DICT_ARUCO_ORIGINAL', help='aruco marker dictionnary (DICT_4X4_50, DICT_4X4_100, DICT_4X4_250, DICT_4X4_1000, DICT_5X5_50, DICT_5X5_100, DICT_5X5_250, DICT_5X5_1000, DICT_6X6_50, DICT_6X6_100, DICT_6X6_250, DICT_6X6_1000, DICT_7X7_50, DICT_7X7_100, DICT_7X7_250, DICT_7X7_1000, DICT_ARUCO_ORIGINAL,DICT_APRILTAG_16h5, DICT_APRILTAG_25h9, DICT_APRILTAG_36h10, DICT_APRILTAG_36h11)')
+    parser.add_argument('-ms', '--marker_size', metavar='MARKER_SIZE', type=float, default=6, help='aruco marker size (cm)')
+    parser.add_argument('-mi', '--marker_id_scene', metavar='MARKER_ID_SCENE', type=json.loads, help='{"marker": "aoi scene filepath"} dictionary')
+    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 markers id to track
+        if args.marker_id_scene == None:
+            print(f'Track any Aruco markers from the {args.marker_dictionary} dictionary')
+        else:
+            print(f'Track Aruco markers {list(args.marker_id_scene.keys())} from the {args.marker_dictionary} dictionary')
+
+        # 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
+                timerange_path = f'[{int(args.time_range[0])}s - {int(args.time_range[1])}s]'
+
+                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.get_duration()/1e6} s\n\twidth: {tobii_segment_video.get_width()} px\n\theight: {tobii_segment_video.get_height()} px')
+
+        # Create aruco camera
+        aruco_camera = ArUcoCamera.ArUcoCamera()
+
+        # Load calibration file
+        if args.camera_calibration != None:
+
+            aruco_camera.load_calibration_file(args.camera_calibration)
+
+        else:
+
+            raise ValueError('.json camera calibration filepath required. Use -c option.')
+
+        # Create aruco tracker
+        aruco_tracker = ArUcoTracker.ArUcoTracker(args.marker_dictionary, args.marker_size, aruco_camera)
+
+        # Load specific configuration file
+        def load_configuration_file():
+
+            if args.aruco_tracker_configuration != None:
+
+                aruco_tracker.load_configuration_file(args.aruco_tracker_configuration)
+
+                print(f'ArUcoTracker configuration for {aruco_tracker.get_markers_dictionay().get_markers_format()} markers detection:')
+                aruco_tracker.print_configuration()
+
+        load_configuration_file()
+
+        # Load AOI 3D scene for each marker
+        aoi3D_scenes = {}
+    
+        for marker_id, aoi_scene_filepath in args.marker_id_scene.items():
+
+            marker_id = int(marker_id)
+            
+            aoi3D_scenes[marker_id] = AOI3DScene.AOI3DScene()
+            aoi3D_scenes[marker_id].load(aoi_scene_filepath)
+
+            print(f'AOI in {os.path.basename(aoi_scene_filepath)} scene related to marker #{marker_id}:')
+            for aoi in aoi3D_scenes[marker_id].keys():
+                print(f'\t{aoi}')
+
+        def aoi3D_scene_selector(marker_id):
+            return aoi3D_scenes.get(marker_id, None)
+
+        # Display first frame
+        video_ts, video_frame = tobii_segment_video.get_frame(0)
+        cv.imshow(f'Segment {tobii_segment.get_id()} ArUco marker editor', video_frame.matrix)
+
+        # Init pointer and click
+        pointer = (0, 0)
+        click = (0, 0)
+
+        # On mouse left click : update pointer position
+        def on_mouse_event(event, x, y, flags, param):
+
+            nonlocal pointer
+            nonlocal click
+
+            # Update pointer
+            pointer = (x, y)
+
+            # Update click
+            if event == cv.EVENT_LBUTTONUP:
+
+                click = pointer
+
+        cv.setMouseCallback(f'Segment {tobii_segment.get_id()} ArUco marker editor', on_mouse_event)
+
+        # Frame selector loop
+        frame_index = 0
+        last_frame_index = -1
+        last_frame_matrix = video_frame.matrix.copy()
+
+        selected_marker_id = -1
+
+        try:
+
+            while True:
+
+                # Select a frame on change
+                if frame_index != last_frame_index:
+
+                    video_ts, video_frame = tobii_segment_video.get_frame(frame_index)
+                    video_ts_ms = video_ts / 1000
+
+                    last_frame_index = frame_index
+                    last_frame_matrix = video_frame.matrix.copy()
+
+                else:
+
+                    video_frame.matrix = last_frame_matrix.copy()
+
+                # Track markers with pose estimation and draw them
+                aruco_tracker.track(video_frame.matrix)
+                aruco_tracker.draw(video_frame.matrix)
+
+                # Write segment timing
+                cv.putText(video_frame.matrix, f'Segment time: {int(video_ts_ms)} ms', (20, 40), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv.LINE_AA)
+
+                # Draw focus area
+                cv.circle(video_frame.matrix, (int(video_frame.width/2), int(video_frame.height/2)), int(video_frame.width/3), (255, 150, 150), 1)
+                     
+                # Draw pointer
+                cv.circle(video_frame.matrix, pointer, 2, (0, 255, 255), -1)
+
+                # Write selected marker id
+                if selected_marker_id >= 0:
+                    cv.putText(video_frame.matrix, f'Marker {selected_marker_id} selected', (20, 80), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv.LINE_AA)
+
+                # Project 3D scene on each video frame and the visualisation frame
+                if aruco_tracker.get_markers_number():
+
+                    # Write detected marker ids
+                    cv.putText(video_frame.matrix, f'Detected markers : {aruco_tracker.get_markers_ids()}', (20, video_frame.height - 40), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv.LINE_AA)
+                            
+                    # Update selected marker id by clicking on marker
+                    for (i, marker_id) in enumerate(aruco_tracker.get_markers_ids()):
+
+                        marker_aoi = numpy.array(aruco_tracker.get_marker_corners(i)).view(AOIFeatures.AreaOfInterest)
+
+                        if marker_aoi.looked(click):
+
+                            selected_marker_id = marker_id
+
+                            # Select 3D scene related to selected marker
+                            aoi3D_scene = aoi3D_scene_selector(selected_marker_id)
+
+                    # If a marker is selected
+                    try:
+                        
+                        # Retreive marker index
+                        selected_marker_index = aruco_tracker.get_marker_index(selected_marker_id)
+
+                        # If AOI scene is found
+                        if aoi3D_scene != None:
+
+                            # Is the marker out of focus area ?
+                            marker_x, marker_y = aruco_tracker.get_marker_center(selected_marker_index)
+                            distance_to_center = ( (video_frame.width/2 - marker_x)**2 + (video_frame.height/2 - marker_y)**2 )**0.5
+
+                            if distance_to_center > int(video_frame.width/3):
+
+                                # Write warning
+                                cv.putText(video_frame.matrix, f'Out of focus area', (20, 120), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv.LINE_AA)
+                            
+                            aoi3D_scene.rotation = aruco_tracker.get_marker_rotation(selected_marker_index)
+                            aoi3D_scene.translation = aruco_tracker.get_marker_translation(selected_marker_index)
+
+                            # Write rotation matrix
+                            R, _ = cv.Rodrigues(aoi3D_scene.rotation)
+                            cv.putText(video_frame.matrix, f'Rotation matrix:', (20, 160), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv.LINE_AA)
+                            cv.putText(video_frame.matrix, f'{R[0][0]:.3f}   {R[0][1]:.3f}   {R[0][2]:.3f}', (40, 200), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv.LINE_AA)
+                            cv.putText(video_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(video_frame.matrix, f'{R[2][0]:.3f}   {R[2][1]:.3f}   {R[2][2]:.3f}', (40, 280), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv.LINE_AA)
+                            
+                            # Write translation vector
+                            T = aoi3D_scene.translation
+                            cv.putText(video_frame.matrix, f'Translation vector:', (20, 320), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv.LINE_AA)
+                            cv.putText(video_frame.matrix, f'{T[0][0]:.3f}', (40, 360), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv.LINE_AA)
+                            cv.putText(video_frame.matrix, f'{T[0][1]:.3f}', (40, 400), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv.LINE_AA)
+                            cv.putText(video_frame.matrix, f'{T[0][2]:.3f}', (40, 440), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv.LINE_AA)
+                            
+                            # Remove aoi outside vision field 
+                            # The vision cone tip is positionned behind the head
+                            #aoi3D_scene = aoi3D_scene.clip(300, 150, cone_tip=[0., 0., -20.])
+
+                            # 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.
+                            aoi2D_video_scene = aoi3D_scene.project(aruco_camera.get_K())
+
+                            # Draw scene
+                            aoi2D_video_scene.draw(video_frame.matrix, pointer, 2, exclude=['Visualisation_Plan'])
+
+                        else:
+
+                            # Write error
+                            cv.putText(video_frame.matrix, f'Marker {selected_marker_id} have no AOI scene', (20, 120), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv.LINE_AA)
+
+                    except ValueError:
+                        
+                        # Write error
+                        if selected_marker_id >= 0:
+                            cv.putText(video_frame.matrix, f'Marker {selected_marker_id} not found', (20, 120), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv.LINE_AA)
+
+                # Reset click
+                click = (0, 0)
+
+                if args.window:
+
+                    key_pressed = cv.waitKey(1)
+
+                    # 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
+
+                    # Close window using 'Esc' key
+                    if key_pressed == 27:
+                        break
+
+                    # Reload tracker configuratio on 't' key
+                    if key_pressed == 116:
+                        load_configuration_file()
+
+                    # Display video
+                    cv.imshow(f'Segment {tobii_segment.get_id()} ArUco marker editor', video_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