1 files changed, 40 insertions, 2 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()
+