Exporting video with gaze data.

1 files changed, 12 insertions, 1 deletions

diff --git a/src/argaze/utils/replay_tobii_session.py b/src/argaze/utils/replay_tobii_session.py
index 08a0ce3..d001889 100644
--- a/src/argaze/utils/replay_tobii_session.py
+++ b/src/argaze/utils/replay_tobii_session.py
@@ -41,6 +41,10 @@ def main():
         tobii_ts_pupil_diameter = tobii_segment_data.gidx_pd_eye
         print(f'{len(tobii_ts_pupil_diameter)} pupil diameters loaded')
 
+        # Prepare video exportation at the same format than segment video
+        gaze_video_filepath = f'{args.segment_path}/gaze.mp4'
+        output_video = TobiiVideo.TobiiVideoOutput(gaze_video_filepath, tobii_segment_video.get_stream())
+
         # Video and data replay loop
         try:
 
@@ -67,7 +71,10 @@ def main():
 
                     # Draw pupil diameter
                     pupil_diameter = int(closest_pupil_diameter.pd * 2)
-                    cv.circle(video_frame.matrix, gaze_position.as_tuple(), pupil_diameter, (255, 0, 255), 1)
+                    cv.circle(video_frame.matrix, gaze_position, pupil_diameter, (255, 0, 255), 1)
+
+                    # Write video
+                    output_video.write(video_frame.matrix)
 
                 # Ignore missing pupil diameter
                 except ValueError:
@@ -86,6 +93,10 @@ def main():
         # Stop frame display
         cv.destroyAllWindows()
 
+        # End output video file
+        output_video.close()
+        print(f'\nVideo with gaze data saved into {gaze_video_filepath}')
+
 if __name__ == '__main__':
 
     main()
\ No newline at end of file