Editing scene rotation and saving it into the obj scene file.

1 files changed, 103 insertions, 31 deletions

diff --git a/src/argaze/utils/edit_tobii_segment_aruco_pose.py b/src/argaze/utils/edit_tobii_segment_aruco_pose.py
index 8748b83..05da07d 100644
--- a/src/argaze/utils/edit_tobii_segment_aruco_pose.py
+++ b/src/argaze/utils/edit_tobii_segment_aruco_pose.py
@@ -104,6 +104,7 @@ def main():
 
         # Load AOI 3D scene for each marker
         aoi3D_scenes = {}
+        aoi3D_scene_edits = {}
     
         for marker_id, aoi_scene_filepath in args.marker_id_scene.items():
 
@@ -112,6 +113,11 @@ def main():
             aoi3D_scenes[marker_id] = AOI3DScene.AOI3DScene()
             aoi3D_scenes[marker_id].load(aoi_scene_filepath)
 
+            aoi3D_scene_edits[marker_id] = {
+                'rotation': numpy.array([0.0, 0.0, 0.0]),
+                'translation': numpy.array([0.0, 0.0, 0.0])
+            }
+
             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}')
@@ -119,27 +125,49 @@ def main():
         def aoi3D_scene_selector(marker_id):
             return aoi3D_scenes.get(marker_id, None)
 
+        def aoi3D_scene_edit_selector(marker_id):
+            return aoi3D_scene_edits.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
+        # Init mouse interaction variables
         pointer = (0, 0)
-        click = (0, 0)
+        left_click = (0, 0)
+        right_click = (0, 0)
+        right_button = False
+        edit_coord = 0 # x
 
-        # On mouse left click : update pointer position
+        # On mouse left left_click : update pointer position
         def on_mouse_event(event, x, y, flags, param):
 
             nonlocal pointer
-            nonlocal click
+            nonlocal left_click
+            nonlocal right_click
+            nonlocal right_button
 
             # Update pointer
             pointer = (x, y)
 
-            # Update click
+            # Update left_click
             if event == cv.EVENT_LBUTTONUP:
 
-                click = pointer
+                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.get_id()} ArUco marker editor', on_mouse_event)
 
@@ -186,7 +214,15 @@ def main():
 
                 # 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)
+
+                    cv.putText(video_frame.matrix, f'Marker {selected_marker_id} : Axis {edit_coord + 1} selected', (20, 80), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv.LINE_AA)
+
+                # Write documentation
+                else: 
+                    cv.putText(video_frame.matrix, f'Left click on marker to select scene', (20, 80), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv.LINE_AA)
+                    cv.putText(video_frame.matrix, f'Shift+num to select axis', (20, 120), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv.LINE_AA)
+                    cv.putText(video_frame.matrix, f'Right click and drag to edit axis', (20, 160), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv.LINE_AA)
+                    cv.putText(video_frame.matrix, f'Ctrl+s to save scene', (20, 200), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv.LINE_AA)
 
                 # Project 3D scene on each video frame and the visualisation frame
                 if aruco_tracker.get_markers_number():
@@ -194,12 +230,12 @@ def main():
                     # 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
+                    # Update selected marker id by left_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):
+                        if marker_aoi.looked(left_click):
 
                             selected_marker_id = marker_id
 
@@ -223,31 +259,51 @@ def main():
 
                                 # 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)
+
+                            # Select scene edit
+                            aoi3D_scene_edit = aoi3D_scene_edit_selector(selected_marker_id)
+
+                            # Edit scene
+                            if aoi3D_scene_edit != None:
+
+                                if right_button:
+
+                                    pointer_delta_x, pointer_delta_y = (right_click[0] - video_frame.width/2) / (video_frame.width/3), (video_frame.height/2 - right_click[1]) / (video_frame.width/3)
+                                    
+                                    # 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], pointer_delta_x])
+
+                                # Apply transformation
+                                aoi3D_scene_edited = aoi3D_scene.transform(aoi3D_scene_edit['translation'], aoi3D_scene_edit['rotation'])
+
+                                # Write rotation matrix
+                                R, _ = cv.Rodrigues(aoi3D_scene_edit['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_edit['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]:.3f}', (40, 360), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 1, cv.LINE_AA)
+                                cv.putText(video_frame.matrix, f'{T[1]:.3f}', (40, 400), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv.LINE_AA)
+                                cv.putText(video_frame.matrix, f'{T[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.])
+                            #aoi3D_scene_edited = aoi3D_scene_edited.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())
+                            aoi2D_video_scene = aoi3D_scene_edited.project(aruco_tracker.get_marker_translation(selected_marker_index), aruco_tracker.get_marker_rotation(selected_marker_index), aruco_camera.get_K())
 
                             # Draw scene
                             aoi2D_video_scene.draw(video_frame.matrix, pointer, 2, exclude=['Visualisation_Plan'])
@@ -263,13 +319,16 @@ def main():
                         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)
+                # 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
@@ -282,6 +341,20 @@ def main():
                     if frame_index < 0:
                         frame_index = 0
 
+                    # Select coordinate to edit
+                    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}']
+                            aoi3D_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
@@ -303,7 +376,6 @@ def main():
         # Stop frame display
         cv.destroyAllWindows()
 
-
 if __name__ == '__main__':
 
     main()
\ No newline at end of file