diff options
| -rw-r--r-- | src/argaze/utils/tobii_segment_arscene_edit.py | 365 |
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 |