From 449db6afcad246b7f7a27c2e0cebed99ac3feb85 Mon Sep 17 00:00:00 2001 From: Théo de la Hogue Date: Wed, 7 Sep 2022 22:40:04 +0200 Subject: Adding a new utils script to edit aruco pose. --- src/argaze/utils/edit_tobii_segment_aruco_pose.py | 305 ++++++++++++++++++++++ 1 file changed, 305 insertions(+) create mode 100644 src/argaze/utils/edit_tobii_segment_aruco_pose.py 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 -- cgit v1.1