aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--src/argaze/utils/tobii_segment_arscene_edit.py365
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