Improving aoi export

1 files changed, 20 insertions, 18 deletions

diff --git a/src/argaze/utils/tobii_segment_aruco_aoi_export.py b/src/argaze/utils/tobii_segment_aruco_aoi_export.py
index bb7d769..57ee916 100644
--- a/src/argaze/utils/tobii_segment_aruco_aoi_export.py
+++ b/src/argaze/utils/tobii_segment_aruco_aoi_export.py
@@ -70,8 +70,8 @@ def main():
                     os.makedirs(destination_path)
                     print(f'{destination_path} folder created')
 
-        vs_data_filepath = f'{destination_path}/aruco_aoi.csv'
-        vs_video_filepath = f'{destination_path}/aruco_aoi.mp4'
+        vs_data_filepath = f'{destination_path}/aoi.csv'
+        vs_video_filepath = f'{destination_path}/aoi.mp4'
 
         # 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)
@@ -150,7 +150,7 @@ def main():
         try:
 
             # Initialise progress bar
-            MiscFeatures.printProgressBar(0, tobii_segment_video.duration/1e3, prefix = 'Progress:', suffix = 'Complete', length = 100)
+            #MiscFeatures.printProgressBar(0, tobii_segment_video.duration/1e3, prefix = 'Progress:', suffix = 'Complete', length = 100)
 
             head_moving = False
             head_movement_last = 0.
@@ -169,6 +169,12 @@ def main():
                     # Get nearest video timestamp
                     _, nearest_vts = tobii_ts_vts.get_last_before(video_ts)
 
+                    # Edit dictionary to store 2D aoi with primary timestamp offset and warning
+                    all_aoi2D = {
+                        'offset': nearest_vts.offset,
+                        'warning': None
+                    }
+
                     # Get nearest head rotation before video timestamp and remove all head rotations before
                     _, nearest_head_rotation = tobii_ts_head_rotations.pop_first_until(video_ts)
 
@@ -191,7 +197,9 @@ def main():
                     # When head is moving, ArUco tracking could return bad pose estimation and so bad AOI scene projection
                     if head_moving:
 
-                        ts_offset_aois[round(video_ts_ms)] = AOIFeatures.EmptyAOIScene()
+                        all_aoi2D['warning'] = 'Head is moving'
+
+                        ts_offset_aois[video_ts] = all_aoi2D
                         
                         raise UserWarning('Head is moving')
 
@@ -206,7 +214,9 @@ def main():
                     # When no marker is detected, no AOI scene projection can't be done
                     if len(aruco_tracker.tracked_markers) == 0:
 
-                        ts_offset_aois[round(video_ts_ms)] = AOIFeatures.EmptyAOIScene()
+                        all_aoi2D['warning'] = 'No marker detected'
+
+                        ts_offset_aois[video_ts] = all_aoi2D
 
                         raise UserWarning('No marker detected')
 
@@ -253,19 +263,11 @@ def main():
 
                     aoi2D_merged_scene.draw(visu_frame.matrix, (0, 0))
 
-                    # Edit dictionary with all 2D aoi with primary timestamp offset
-                    all_aoi2D = {
-                        'offset': nearest_vts.offset
-                    }
-
-                    for aoi_name in all_aois_names:
-
-                        try:
-                            all_aoi2D[aoi_name] = aoi2D_merged_scene[aoi_name]
-                        except:
-                            all_aoi2D[aoi_name] = None
-
                     # Store all 2D aoi
+                    for aoi_name in aoi2D_merged_scene.keys():
+
+                        all_aoi2D[aoi_name] = numpy.rint(aoi2D_merged_scene[aoi_name]).astype(int)
+ 
                     ts_offset_aois[video_ts] = all_aoi2D
 
                     # Warn user when the merged scene is empty
@@ -308,7 +310,7 @@ def main():
 
                 # Update Progress Bar
                 progress = video_ts_ms - int(args.time_range[0] * 1e3)
-                MiscFeatures.printProgressBar(progress, tobii_segment_video.duration/1e3, prefix = 'Progress:', suffix = 'Complete', length = 100)
+                #MiscFeatures.printProgressBar(progress, tobii_segment_video.duration/1e3, prefix = 'Progress:', suffix = 'Complete', length = 100)
 
         # Exit on 'ctrl+C' interruption
         except KeyboardInterrupt: