#!/usr/bin/env python import argparse import time import itertools from argaze import ArFeatures, GazeFeatures from argaze.AreaOfInterest import AOIFeatures from argaze.ArUcoMarkers import ArUcoScene from argaze.utils import MiscFeatures from tobiiproglasses2 import * import cv2 import numpy def main(): """ Load AR environment from .json file, detect ArUco markers into movie frames and estimate environment pose. Edit environment setup to improve pose estimation. """ # Manage arguments parser = argparse.ArgumentParser(description=main.__doc__.split('-')[0]) parser.add_argument('environment', metavar='ENVIRONMENT', type=str, help='ar environment filepath') parser.add_argument('movie', metavar='MOVIE', type=str, default=None, help='movie path') parser.add_argument('-s','--start', metavar='START', type=float, default=0., help='start time in second') parser.add_argument('-o', '--output', metavar='OUT', type=str, default='environment.json', help='edited ar environment file path') args = parser.parse_args() # Load AR enviroment ar_environment = ArFeatures.ArEnvironment.from_json(args.environment) #print('ArEnvironment:\n', ar_environment) # Select first AR scene ar_scene = list(ar_environment.scenes.values())[0] # Create a window to display AR environment cv2.namedWindow(ar_environment.name, cv2.WINDOW_AUTOSIZE) # Init mouse interaction pointer = (0, 0) left_click = (0, 0) right_click = (0, 0) right_drag = (0, 0) right_button = False edit_trans = False # translate edit_z = False draw_help = False draw_grid = False draw_cover = False pose_mode = 0 # Update pointer position def on_mouse_event(event, x, y, flags, param): nonlocal pointer nonlocal left_click nonlocal right_click nonlocal right_drag nonlocal right_button # Update pointer pointer = (x, y) # Update left_click if event == cv2.EVENT_LBUTTONUP: left_click = pointer # Udpate right_button elif event == cv2.EVENT_RBUTTONDOWN and not right_button: right_button = True right_click = pointer elif event == cv2.EVENT_RBUTTONUP and right_button: right_button = False # Udpate right_drag if right_button: right_drag = (pointer[0] - right_click[0], pointer[1] - right_click[1]) # Attach mouse callback to window cv2.setMouseCallback(ar_environment.name, on_mouse_event) # Enable movie video capture video_capture = cv2.VideoCapture(args.movie) video_fps = video_capture.get(cv2.CAP_PROP_FPS) frame_width = int(video_capture.get(cv2.CAP_PROP_FRAME_WIDTH)) frame_height = int(video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) # Enable exit signal handler exit = MiscFeatures.ExitSignalHandler() # Init frame selection current_frame_index = -1 _, current_frame = video_capture.read() next_frame_index = int(args.start * video_fps) refresh = False # Init marker selection selected_marker_id = -1 # Init place edition place_edit = {} while not exit.status(): # Edit fake gaze position from pointer gaze_position = GazeFeatures.GazePosition(pointer, precision=2) # Reset info frame info_frame = numpy.full((500, 1000, 3), 0, dtype=numpy.uint8) # Select a new frame and detect markers once if next_frame_index != current_frame_index or refresh or draw_cover: video_capture.set(cv2.CAP_PROP_POS_FRAMES, next_frame_index) success, video_frame = video_capture.read() if success: # Refresh once refresh = False current_frame_index = video_capture.get(cv2.CAP_PROP_POS_FRAMES) - 1 current_frame_time = video_capture.get(cv2.CAP_PROP_POS_MSEC) # Draw camera calibration if draw_grid: ar_environment.aruco_detector.camera.draw(video_frame, frame_width/10, frame_height/10, 100., color=(127, 127, 127)) # Hide zone if draw_cover: # Draw black circle under pointer cv2.circle(video_frame, pointer, 50, (0, 0, 0), -1) # Detect markers ar_environment.aruco_detector.detect_markers(video_frame) # Edit marker's color for i, m in ar_environment.aruco_detector.detected_markers.items(): m.color = list(itertools.permutations([0, 255, 255]))[i] # Draw center cv2.circle(video_frame, m.center.astype(int), 5, m.color, -1) # Write timing cv2.rectangle(video_frame, (0, 0), (frame_width, 50), (63, 63, 63), -1) cv2.putText(video_frame, f'Time: {int(current_frame_time)} ms', (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA) # Copy frame current_frame = video_frame.copy() # Keep last frame else: video_frame = current_frame.copy() # Handle marker selection on left click if len(ar_environment.aruco_detector.detected_markers) > 0: # Update selected marker id by left clicking on marker for (marker_id, marker) in ar_environment.aruco_detector.detected_markers.items(): marker_aoi = marker.corners.reshape(4, 2).view(AOIFeatures.AreaOfInterest) if marker_aoi.contains_point(left_click): selected_marker_id = marker_id break else: selected_marker_id = -1 try: # A marker is selected if selected_marker_id >= 0: try: # Retreive selected marker selected_marker = ar_environment.aruco_detector.detected_markers[selected_marker_id] # Write selected marker id cv2.rectangle(info_frame, (0, 0), (1000, 50), selected_marker.color, -1) cv2.putText(info_frame, f'Selected marker #{selected_marker.identifier}', (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 1, cv2.LINE_AA) # Estimate selected marker pose ar_environment.aruco_detector.estimate_markers_pose([selected_marker_id]) # Write selected marker rotation matrix R = selected_marker.rotation cv2.putText(info_frame, f'Rotation matrix:', (20, 120), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA) cv2.putText(info_frame, f'{R[0][0]:.3f} {R[0][1]:.3f} {R[0][2]:.3f}', (40, 160), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA) cv2.putText(info_frame, f'{R[1][0]:.3f} {R[1][1]:.3f} {R[1][2]:.3f}', (40, 200), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv2.LINE_AA) cv2.putText(info_frame, f'{R[2][0]:.3f} {R[2][1]:.3f} {R[2][2]:.3f}', (40, 240), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv2.LINE_AA) # Write selected marker translation vector T = selected_marker.translation cv2.putText(info_frame, f'Translation vector:', (20, 320), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA) cv2.putText(info_frame, f'{T[0]:.3f}', (40, 360), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA) cv2.putText(info_frame, f'{T[1]:.3f}', (40, 400), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv2.LINE_AA) cv2.putText(info_frame, f'{T[2]:.3f}', (40, 440), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv2.LINE_AA) # Retreive selected marker place selected_place = ar_scene.aruco_scene.places[selected_marker_id] cv2.putText(info_frame, f'Edited place #{selected_place.marker.identifier}', (520, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 1, cv2.LINE_AA) # On right click if right_button: pointer_delta_x, pointer_delta_y = right_drag[0] / frame_width, right_drag[1] / frame_height place_edit[selected_marker_id] = {'rotation': (0, 0, 0), 'translation': (0, 0, 0)} if edit_trans: # Edit place rotation if edit_z: place_edit[selected_marker_id]['rotation'] = (0, 0, -pointer_delta_y) else: place_edit[selected_marker_id]['rotation'] = (pointer_delta_y, pointer_delta_x, 0) else: # Edit place translation if edit_z: place_edit[selected_marker_id]['translation'] = (0, 0, pointer_delta_y) else: place_edit[selected_marker_id]['translation'] = (-pointer_delta_x, pointer_delta_y, 0) # Apply transformations R = selected_place.rotation.dot(ArUcoScene.make_rotation_matrix(*place_edit[selected_marker_id]['rotation']).T) T = selected_place.translation + numpy.array(place_edit[selected_marker_id]['translation']) edited_place = ArUcoScene.Place(T, R, selected_marker) else: edited_place = selected_place # Write edited place rotation matrix R = edited_place.rotation cv2.putText(info_frame, f'Rotation matrix:', (520, 120), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA) cv2.putText(info_frame, f'{R[0][0]:.3f} {R[0][1]:.3f} {R[0][2]:.3f}', (540, 160), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA) cv2.putText(info_frame, f'{R[1][0]:.3f} {R[1][1]:.3f} {R[1][2]:.3f}', (540, 200), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv2.LINE_AA) cv2.putText(info_frame, f'{R[2][0]:.3f} {R[2][1]:.3f} {R[2][2]:.3f}', (540, 240), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv2.LINE_AA) # Write edited place translation vector T = edited_place.translation cv2.putText(info_frame, f'Translation vector:', (520, 320), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA) cv2.putText(info_frame, f'{T[0]:.3f}', (540, 360), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA) cv2.putText(info_frame, f'{T[1]:.3f}', (540, 400), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv2.LINE_AA) cv2.putText(info_frame, f'{T[2]:.3f}', (540, 440), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv2.LINE_AA) # Replace selected place by edited place ar_scene.aruco_scene.places[selected_marker_id] = edited_place # Estimate scene pose considering only selected marker tvec, rmat, strategy, _ = ar_scene.estimate_pose({selected_marker_id: selected_marker}) strategy = strategy.replace('_', ' ') # Write pose estimation strategy cv2.putText(video_frame, f'{strategy} with marker {selected_marker_id}', (20, frame_height - 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA) # Draw expected marker places ar_scene.draw_places(video_frame) # Project AOI scene into frame according estimated pose aoi_scene_projection = ar_scene.project(tvec, rmat, visual_hfov=TobiiSpecifications.VISUAL_HFOV) # Draw AOI scene projection with gaze aoi_scene_projection.draw_circlecast(video_frame, gaze_position, base_color=selected_marker.color, looked_color=selected_marker.color) # Catch missing selected marker except KeyError: cv2.putText(video_frame, f'Marker {selected_marker_id} not found', (20, 120), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA) # No marker selected else: cv2.putText(info_frame, f'Left click on marker to select it', (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA) # Estimate all marker's pose ar_environment.aruco_detector.estimate_markers_pose() # Check markers consistency consistent_markers, unconsistent_markers, unconsistencies = ar_scene.aruco_scene.check_markers_consistency(ar_environment.aruco_detector.detected_markers, ar_scene.angle_tolerance, ar_scene.distance_tolerance) # Set unconsistent marker color to red for i, m in ar_environment.aruco_detector.detected_markers.items(): if i in list(unconsistent_markers.keys()): m.color = (0, 0, 255) # Write unconsistencies for i, (label, value) in enumerate(unconsistencies.items()): cv2.putText(info_frame, f'Unconsistent {label}: {value:.3f}', (20, 120+ i*40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA) # Single marker scene pose estimation if pose_mode == 0: cv2.putText(video_frame, f'Single marker scene pose estimation', (20, frame_height - 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA) for i, m in ar_environment.aruco_detector.detected_markers.items(): tvec, rmat = ar_scene.aruco_scene.estimate_pose_from_single_marker(m) # Project AOI scene into frame according estimated pose aoi_scene_projection = ar_scene.project(tvec, rmat, visual_hfov=TobiiSpecifications.VISUAL_HFOV) aoi_scene_projection.draw(video_frame, color=m.color) # ArUco marker axis scene pose estimation elif pose_mode == 1: # Write pose estimation strategy cv2.putText(video_frame, f'ArUco marker axis scene pose estimation', (20, frame_height - 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA) for axis_name, axis_markers in ar_scene.aruco_axis.items(): try: origin_marker = ar_environment.aruco_detector.detected_markers[axis_markers['origin_marker']] horizontal_axis_marker = ar_environment.aruco_detector.detected_markers[axis_markers['horizontal_axis_marker']] vertical_axis_marker = ar_environment.aruco_detector.detected_markers[axis_markers['vertical_axis_marker']] tvec, rmat = ar_scene.aruco_scene.estimate_pose_from_axis_markers(origin_marker, horizontal_axis_marker, vertical_axis_marker) # Project AOI scene into frame according estimated pose aoi_scene_projection = ar_scene.project(tvec, rmat, visual_hfov=TobiiSpecifications.VISUAL_HFOV) aoi_scene_projection.draw(video_frame, color=(255, 255, 255)) break except: pass # ArUco AOI scene building elif pose_mode == 2: # Write pose estimation strategy cv2.putText(video_frame, f'ArUco AOI scene building', (20, frame_height - 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA) try : # Try to build AOI scene from detected ArUco marker corners aoi_scene_projection = ar_scene.build_aruco_aoi_scene(ar_environment.aruco_detector.detected_markers) except: pass # Draw expected marker places #ar_scene.draw_places(video_frame) # Catch exceptions raised by estimate_pose and project methods except (ArFeatures.PoseEstimationFailed, ArFeatures.SceneProjectionFailed) as e: cv2.rectangle(video_frame, (0, 90), (700, 130), (127, 127, 127), -1) cv2.putText(video_frame, f'Error: {e}', (20, 100), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA) # Draw frame cv2.imshow(ar_environment.name, video_frame) # Draw detected markers ar_environment.aruco_detector.draw_detected_markers(video_frame) # Draw pointer gaze_position.draw(video_frame) # Write documentation cv2.putText(video_frame, f'Press \'h\' for help', (950, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA) if draw_help: cv2.rectangle(video_frame, (0, 50), (700, 300), (127, 127, 127), -1) cv2.putText(video_frame, f'> Left click on marker: select marker', (20, 80), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA) cv2.putText(video_frame, f'> Left click on frame: unselect marker', (20, 120), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA) cv2.putText(video_frame, f'> T: translate, R: rotate, Z: select axis', (20, 160), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA) cv2.putText(video_frame, f'> Right click and drag: edit axis', (20, 200), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA) cv2.putText(video_frame, f'> Ctrl + S: save environment', (20, 240), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA) cv2.putText(video_frame, f'> Backspace: reload environment', (20, 280), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA) # Write selected marker id elif selected_marker_id >= 0: cv2.rectangle(video_frame, (0, 50), (700, 90), (127, 127, 127), -1) # Select color if edit_z: str_axis = 'Z' color_axis = (255, 0, 0) else: str_axis = 'XY' color_axis = (0, 255, 255) if edit_trans: cv2.putText(video_frame, f'Rotate marker {selected_marker_id} around axis {str_axis}', (20, 80), cv2.FONT_HERSHEY_SIMPLEX, 1, color_axis, 1, cv2.LINE_AA) else: cv2.putText(video_frame, f'Translate marker {selected_marker_id} along axis {str_axis}', (20, 80), cv2.FONT_HERSHEY_SIMPLEX, 1, color_axis, 1, cv2.LINE_AA) key_pressed = cv2.waitKey(10) #if key_pressed != -1: # print(key_pressed) # Select previous frame with left arrow if key_pressed == 2: next_frame_index -= 1 # Select next frame with right arrow if key_pressed == 3: next_frame_index += 1 # Clip frame index if next_frame_index < 0: next_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 # Switch Z axis edition if key_pressed == 122: edit_z = not edit_z # Switch help mode with h key if key_pressed == 104: draw_help = not draw_help # Switch grid mode with g key if key_pressed == 103: draw_grid = not draw_grid refresh = True # Switch draw_cover mode with c key if key_pressed == 99: draw_cover = not draw_cover # Switch pose estimation mode with m key if key_pressed == 109: pose_mode += 1 if pose_mode > 2: pose_mode = 0 # Save selected marker edition using 'Ctrl + s' if key_pressed == 19: ar_environment.to_json(args.output) print(f'Environment saved into {args.output}') # Close window using 'Esc' key if key_pressed == 27: break # Reload environment on 'Backspace' key if key_pressed == 127: ar_environment = ArFeatures.ArEnvironment.from_json(args.environment) print(f'Environment reloaded from {args.environment}') refresh = True # Display video cv2.imshow(ar_environment.name, video_frame) # Display info cv2.imshow('Info', info_frame) # Close movie capture video_capture.release() # Stop frame display cv2.destroyAllWindows() if __name__ == '__main__': main()