1 files changed, 30 insertions, 6 deletions
diff --git a/src/argaze/TobiiGlassesPro2/TobiiVideo.py b/src/argaze/TobiiGlassesPro2/TobiiVideo.py
index 6318de8..e6ec064 100644
--- a/src/argaze/TobiiGlassesPro2/TobiiVideo.py
+++ b/src/argaze/TobiiGlassesPro2/TobiiVideo.py
@@ -22,10 +22,11 @@ class TobiiVideoFrame(DataStructures.DictObject):
 class TobiiVideoSegment():
     """Handle Tobii Glasses Pro 2 segment video file."""
 
-    def __init__(self, segment_video_path):
+    def __init__(self, segment_video_path, start_timestamp:int = 0, end_timestamp:int = None):
         """Load segment video from segment directory"""
 
         self.__segment_video_path = segment_video_path
+        
         self.__container = av.open(self.__segment_video_path)
         self.__stream = self.__container.streams.video[0]
 
@@ -33,15 +34,23 @@ class TobiiVideoSegment():
         self.__height = int(cv.VideoCapture(self.__segment_video_path).get(cv.CAP_PROP_FRAME_HEIGHT))
 
         self.__vts_data_buffer = None
+        self.__vts_offset = 0
+
+        self.__start_timestamp = start_timestamp
+        self.__end_timestamp = end_timestamp
+
+        # position at the given start time
+        self.__container.seek(self.__start_timestamp)
         
     def get_path(self):
         return self.__segment_video_path
 
     def get_duration(self):
-        return float(self.__stream.duration * self.__stream.time_base)
-
-    def get_frame_number(self):
-        return self.__stream.frames
+        """Duration in microsecond"""
+        if self.__end_timestamp == None:
+            return int((self.__stream.duration * self.__stream.time_base) * 1000000) - self.__start_timestamp
+        else:
+            return self.__end_timestamp - self.__start_timestamp
 
     def get_width(self):
         return self.__width
@@ -52,8 +61,11 @@ class TobiiVideoSegment():
     def get_stream(self):
         return self.__stream
 
+    def get_vts_offset(self):
+        return self.__vts_offset
+
     def frames(self, vts_data_buffer = None):
-        """Access to frame iterator and optionnaly setup vide / data timestamp synchronisation through vts data buffer."""
+        """Access to frame iterator and optionnaly setup video / data timestamp synchronisation through vts data buffer."""
 
         self.__vts_data_buffer = vts_data_buffer
 
@@ -61,6 +73,12 @@ class TobiiVideoSegment():
         if self.__vts_data_buffer != None:
 
             self.__vts_ts, self.__vts = self.__vts_data_buffer.pop_first()
+            
+            # pop vts buffer until start timestamp
+            while self.__start_timestamp > self.__vts.vts:
+                if len(self.__vts_data_buffer) > 0:
+                    self.__vts_ts, self.__vts = self.__vts_data_buffer.pop_first()
+
             self.__vts_offset = (self.__vts_ts - self.__vts.vts)
 
         return self.__iter__()
@@ -78,6 +96,12 @@ class TobiiVideoSegment():
 
         video_ts = int(frame.time * 1000000)
 
+        # Ignore frames after end timestamp
+        if self.__end_timestamp != None:
+
+            if video_ts >= self.__end_timestamp:
+                raise StopIteration
+
         # If video / data synchronisation is active
         if self.__vts_data_buffer != None: