Adding pose estimation mode. Adding Info window. Unselecting marker by clicking on frame.

1 files changed, 126 insertions, 51 deletions

diff --git a/src/argaze/utils/environment_edit.py b/src/argaze/utils/environment_edit.py
index 0093be7..29fe384 100644
--- a/src/argaze/utils/environment_edit.py
+++ b/src/argaze/utils/environment_edit.py
@@ -2,6 +2,7 @@
 
 import argparse
 import time
+import itertools
 
 from argaze import ArFeatures, GazeFeatures
 from argaze.AreaOfInterest import AOIFeatures
@@ -48,6 +49,7 @@ def main():
     edit_z = False
     cover = False
     draw_help = False
+    pose_mode = 0
 
     # Update pointer position
     def on_mouse_event(event, x, y, flags, param):
@@ -105,12 +107,15 @@ def main():
 
     # 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_detection or cover:
 
@@ -132,17 +137,13 @@ def main():
                 # Detect markers
                 ar_environment.aruco_detector.detect_markers(video_frame)
 
-                # Draw detected markers
-                ar_environment.aruco_detector.draw_detected_markers(video_frame)
-
                 # Draw focus area
                 cv2.rectangle(video_frame, (int(frame_width/6), 0), (int(frame_width*(1-1/6)), int(frame_height)), (255, 150, 150), 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)
-                cv2.putText(video_frame, f'Press \'h\' for help', (950, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)
-
+                
                 # Copy frame
                 current_frame = video_frame.copy()
 
@@ -162,6 +163,11 @@ def main():
                 if marker_aoi.contains_point(left_click):
 
                     selected_marker_id = marker_id
+                    break
+
+                else:
+
+                    selected_marker_id = -1
 
             try:
 
@@ -173,12 +179,31 @@ def main():
                         # Retreive selected marker
                         selected_marker = ar_environment.aruco_detector.detected_markers[selected_marker_id]
 
+                        # Write selected marker id
+                        cv2.putText(info_frame, f'Selected marker #{selected_marker.identifier}', (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 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, (255, 255, 255), 1, cv2.LINE_AA)
+
                         # On right click
                         if right_button:
 
@@ -212,21 +237,20 @@ def main():
 
                             edited_place = selected_place
                         
-                        cv2.rectangle(video_frame, (0, 130), (460, 450), (127, 127, 127), -1)
-                        
-                        # Write edited rotation matrix
+                        # Write edited place rotation matrix
                         R = edited_place.rotation
-                        cv2.putText(video_frame, f'Rotation matrix:', (20, 160), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
-                        cv2.putText(video_frame, f'{R[0][0]:.3f}   {R[0][1]:.3f}   {R[0][2]:.3f}', (40, 200), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
-                        cv2.putText(video_frame, f'{R[1][0]:.3f}   {R[1][1]:.3f}   {R[1][2]:.3f}', (40, 240), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv2.LINE_AA)
-                        cv2.putText(video_frame, f'{R[2][0]:.3f}   {R[2][1]:.3f}   {R[2][2]:.3f}', (40, 280), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv2.LINE_AA)
+                        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 translation vector
+                        # Write edited place translation vector
                         T = edited_place.translation
-                        cv2.putText(video_frame, f'Translation vector:', (20, 320), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
-                        cv2.putText(video_frame, f'{T[0]:.3f}', (40, 360), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
-                        cv2.putText(video_frame, f'{T[1]:.3f}', (40, 400), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv2.LINE_AA)
-                        cv2.putText(video_frame, f'{T[2]:.3f}', (40, 440), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv2.LINE_AA)
+                        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
@@ -255,34 +279,82 @@ def main():
                 # No marker selected
                 else:
 
-                    try:
+                    cv2.putText(info_frame, f'Left click on marker to select it', (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
 
-                        # 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)
+                    # Estimate all marker's pose
+                    ar_environment.aruco_detector.estimate_markers_pose()
 
-                        # Write pose estimation strategy
-                        cv2.putText(video_frame, f'build aruco aoi scene with markers {list(ar_environment.aruco_detector.detected_markers.keys())}', (20, frame_height - 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
+                    # 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)
 
-                    except:
+                    # Edit marker's color
+                    for i, m in ar_environment.aruco_detector.detected_markers.items():
 
-                        # Estimate all marker's pose
-                        ar_environment.aruco_detector.estimate_markers_pose()
+                        m.color = list(itertools.permutations([0, 255, 255]))[i]
 
-                        # Estimate scene pose considering all markers
-                        tvec, rmat, strategy, consistent_markers = ar_scene.estimate_pose(ar_environment.aruco_detector.detected_markers)
-                        strategy = strategy.replace('_', ' ')
+                        # Draw center
+                        cv2.circle(video_frame, m.center.astype(int), 5, m.color, -1)
+
+                        # Set unconsistent marker color to red 
+                        if i in list(unconsistent_markers.keys()):
+                            m.color = (0, 0, 255)
+ 
+                    # 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'{strategy} with markers {list(consistent_markers.keys())}', (20, frame_height - 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
+                        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)
                 
-                        # Draw expected marker places
-                        ar_scene.draw_places(video_frame)
+                        for axis_name, axis_markers in ar_scene.aruco_axis.items():
 
-                        # Project AOI scene into frame according estimated pose
-                        aoi_scene_projection = ar_scene.project(tvec, rmat, visual_hfov=TobiiSpecifications.VISUAL_HFOV)
+                            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)
 
-                    # Draw AOI scene projection with gaze
-                    aoi_scene_projection.draw_circlecast(video_frame, gaze_position)
+                                # 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, 0, 0))
+
+                                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:
@@ -295,20 +367,21 @@ def main():
 
         # 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, 130), (700, 380), (127, 127, 127), -1)
-            cv2.putText(video_frame, f'> Left click on marker: select marker', (20, 160), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)
-            cv2.putText(video_frame, f'> Backspace: unselect marker', (20, 200), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)
-            cv2.putText(video_frame, f'> T: translate, R: rotate', (20, 240), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)
-            cv2.putText(video_frame, f'> Z: switch Z axis edition', (20, 280), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)
-            cv2.putText(video_frame, f'> Right click and drag: edit XY axis', (20, 320), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)
-            cv2.putText(video_frame, f'> Ctrl + S: save scene', (20, 360), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)
+            cv2.rectangle(video_frame, (0, 50), (700, 260), (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)
 
         # Write selected marker id
         elif selected_marker_id >= 0:
@@ -328,9 +401,6 @@ def main():
             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)
 
-        # Reset left_click
-        left_click = (0, 0)
-
         key_pressed = cv2.waitKey(10)
 
         #if key_pressed != -1:
@@ -348,10 +418,6 @@ def main():
         if next_frame_index < 0:
             next_frame_index = 0
 
-        # Unselect marker with backspace key
-        if key_pressed == 127:
-            selected_marker_id = -1
-
         # Edit rotation with r key
         if key_pressed == 114:
             edit_trans = True
@@ -372,6 +438,12 @@ def main():
         if key_pressed == 99:
             cover = not 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)
@@ -389,6 +461,9 @@ def main():
         # Display video
         cv2.imshow(ar_environment.name, video_frame)
 
+        # Display info
+        cv2.imshow('Info', info_frame)
+
     # Close movie capture
     video_capture.release()