Adding TobiiVideoOutput to export data visualisation into a video file

3 files changed, 62 insertions, 7 deletions

diff --git a/src/argaze/TobiiGlassesPro2/TobiiVideo.py b/src/argaze/TobiiGlassesPro2/TobiiVideo.py
index 4d684ec..b3e11f3 100644
--- a/src/argaze/TobiiGlassesPro2/TobiiVideo.py
+++ b/src/argaze/TobiiGlassesPro2/TobiiVideo.py
@@ -49,6 +49,9 @@ class TobiiVideoSegment():
     def get_height(self):
         return self.__height
 
+    def get_stream(self):
+        return self.__stream
+
     def frames(self, vts_data_buffer = None):
         """Access to frame iterator and optionnaly setup vide / data timestamp synchronisation through vts data buffer."""
 
@@ -88,7 +91,7 @@ class TobiiVideoSegment():
             video_ts += self.__vts_offset
 
         # return micro second timestamp and frame data
-        return video_ts, TobiiVideoFrame(frame.to_ndarray(format='bgr24'), frame.width, frame.height)
+        return video_ts, TobiiVideoFrame(frame.to_ndarray(format='rgb24'), frame.width, frame.height)
 
 class TobiiVideoStream(threading.Thread):
     """Capture Tobii Glasses Pro 2 video camera stream."""
@@ -158,7 +161,7 @@ class TobiiVideoStream(threading.Thread):
             self.__read_lock.acquire()
 
             # store frame time, matrix, width, height and pts into a tuple
-            self.__frame_tuple = (frame.time, frame.to_ndarray(format='bgr24'), frame.width, frame.height)
+            self.__frame_tuple = (frame.time, frame.to_ndarray(format='rgb24'), frame.width, frame.height)
 
             # unlock frame access
             self.__read_lock.release()
@@ -179,3 +182,38 @@ class TobiiVideoStream(threading.Thread):
         self.__read_lock.release()
 
         return int(frame_tuple[0] * 1000000), TobiiVideoFrame(frame_tuple[1], frame_tuple[2], frame_tuple[3])
+
+class TobiiVideoOutput():
+    """Export a video file at the same format than a given referent stream."""
+    # TODO : Make a generic video managment to handle video from any device (not only Tobii)
+
+    def __init__(self, output_video_path: str, referent_stream: av.stream.Stream):
+        """Create a video file"""
+
+        self.__output_video_path = output_video_path
+        self.__container = av.open(self.__output_video_path, 'w')
+        self.__stream = self.__container.add_stream(\
+            referent_stream.codec_context.name, \
+            width=referent_stream.codec_context.width, \
+            height=referent_stream.codec_context.height, \
+            rate=referent_stream.codec_context.framerate, \
+            gop_size=referent_stream.codec_context.gop_size, \
+            pix_fmt=referent_stream.codec_context.pix_fmt, \
+            bit_rate=referent_stream.codec_context.bit_rate)
+
+    def get_path(self):
+        return self.__output_video_path
+
+    def write(self, frame):
+        """Write a frame into the output video file"""
+
+        formated_frame = av.VideoFrame.from_ndarray(frame, format='rgb24')
+        formated_frame.reformat(format=self.__stream.codec_context.pix_fmt, interpolation=None)
+        self.__container.mux(self.__stream.encode(formated_frame))
+
+    def close(self):
+        """End the writing of the video file"""
+
+        self.__container.mux(self.__stream.encode()) 
+        self.__container.close()
+
diff --git a/src/argaze/utils/README.md b/src/argaze/utils/README.md
index 99b8717..dc40f30 100644
--- a/src/argaze/utils/README.md
+++ b/src/argaze/utils/README.md
@@ -72,7 +72,7 @@ python ./src/argaze/utils/replay_tobii_session.py -s SEGMENT_PATH
 python ./src/argaze/utils/export_tobii_segment_fixations.py -s SEGMENT_PATH
 ```
 
-- Track ArUco markers into a Tobii camera video segment (replace SEGMENT_PATH). Load an roi scene (replace ROI_SCENE) .obj file, position it virtually relatively to any detected ArUco markers and project the scene into camera frame. Then, detect if Tobii gaze point is inside any ROI.
+- Track ArUco markers into a Tobii camera video segment (replace SEGMENT_PATH). Load an roi scene (replace ROI_SCENE) .obj file, position it virtually relatively to any detected ArUco markers and project the scene into camera frame. Then, detect if Tobii gaze point is inside any ROI. Export ROIs video and data.
 
 ```
 python ./src/argaze/utils/export_tobii_segment_aruco_rois.py -s SEGMENT_PATH -c export/tobii_camera.json -m 7.5 -r ROI_SCENE
diff --git a/src/argaze/utils/export_tobii_segment_aruco_rois.py b/src/argaze/utils/export_tobii_segment_aruco_rois.py
index 343a107..5314141 100644
--- a/src/argaze/utils/export_tobii_segment_aruco_rois.py
+++ b/src/argaze/utils/export_tobii_segment_aruco_rois.py
@@ -15,17 +15,20 @@ import cv2 as cv
 
 def main():
     """
-    Replay Tobii segment video
+    Track any ArUco marker into Tobii Glasses Pro 2 segment video file. 
+    From a loaded ROI scene .obj file, position the scene virtually relatively to any detected ArUco markers and project the scene into camera frame. 
+    Then, detect if Tobii gaze point is inside any ROI.
+    Export ROIs video and data.
     """
 
-    # manage arguments
+    # Manage arguments
     parser = argparse.ArgumentParser(description=main.__doc__.split('-')[0])
     parser.add_argument('-s', '--segment_path', metavar='SEGMENT_PATH', type=str, default=None, help='segment path')
     parser.add_argument('-c', '--camera_calibration', metavar='CAM_CALIB', type=str, default='tobii_camera.json', help='json camera calibration filepath')
     parser.add_argument('-r', '--roi_scene', metavar='ROI_SCENE', type=str, default='roi3D_scene.obj', help='obj roi scene filepath')
     parser.add_argument('-d', '--dictionary', metavar='DICT', type=str, default='DICT_ARUCO_ORIGINAL', help='aruco marker dictionnary')
     parser.add_argument('-m', '--marker_size', metavar='MKR', type=float, default=6, help='aruco marker size (cm)')
-    parser.add_argument('-o', '--output', metavar='OUT', type=str, default=None, help='destination path (segment folder by default)')
+    parser.add_argument('-o', '--output', metavar='OUT', type=str, default=None, help='destination folder path (segment folder by default)')
     args = parser.parse_args()
 
     if args.segment_path != None:
@@ -38,11 +41,13 @@ def main():
                 os.makedirs(os.path.dirname(args.output))
                 print(f'{os.path.dirname(args.output)} folder created')
 
-            rois_filepath = args.output
+            rois_filepath = f'{args.output}/rois.json'
+            video_filepath = f'{args.output}/fullstream+visu.mp4'
 
         else:
 
             rois_filepath = f'{args.segment_path}/rois.json'
+            video_filepath = f'{args.segment_path}/fullstream+visu.mp4'
 
         # Load a tobii segment
         tobii_segment = TobiiEntities.TobiiSegment(args.segment_path)
@@ -59,6 +64,9 @@ def main():
         tobii_ts_gaze_positions = tobii_segment_data.gidx_l_gp
         print(f'{len(tobii_ts_gaze_positions)} gaze positions loaded')
 
+        # Prepare video exportation at the same format than segment video
+        output_video = TobiiVideo.TobiiVideoOutput(video_filepath, tobii_segment_video.get_stream())
+
         # Create aruco camera
         aruco_camera = ArUcoCamera.ArUcoCamera()
         aruco_camera.load_calibration_file(args.camera_calibration)
@@ -132,8 +140,12 @@ def main():
                 if cv.waitKey(1) == 27:
                     break
 
+                # Display video
                 cv.imshow(f'Segment {tobii_segment.get_id()} video', video_frame.matrix)
 
+                # Write video
+                output_video.write(video_frame.matrix)
+
         # Exit on 'ctrl+C' interruption
         except KeyboardInterrupt:
             pass
@@ -141,6 +153,11 @@ def main():
         # Stop frame display
         cv.destroyAllWindows()
 
+        # End output video file
+        output_video.close()
+
+        print(f'ROIs video saved into {video_filepath}')
+
         # Export 2D rois
         roi2D_timestamped_buffer.export_as_json(rois_filepath)