From 148bb63aa8c281b01ce8bd00819af28f8ecbcfd6 Mon Sep 17 00:00:00 2001
From: LinasKo <linkondrackis@gmail.com>
Date: Tue, 13 Aug 2024 16:08:26 +0300
Subject: [PATCH] Revert slow frame seek, add manual seek mode

---
 supervision/utils/video.py | 24 +++++++++++++++++-------
 1 file changed, 17 insertions(+), 7 deletions(-)

diff --git a/supervision/utils/video.py b/supervision/utils/video.py
index 00829ec5..57f5e987 100644
--- a/supervision/utils/video.py
+++ b/supervision/utils/video.py
@@ -118,7 +118,7 @@ class VideoSink:
         self.__writer.release()
 
 
-def _validate_and_setup_video(source_path: str, start: int, end: Optional[int]):
+def _validate_and_setup_video(source_path: str, start: int, end: Optional[int], manual_seek: bool = False):
     video = cv2.VideoCapture(source_path)
     if not video.isOpened():
         raise Exception(f"Could not open video at {source_path}")
@@ -127,11 +127,21 @@ def _validate_and_setup_video(source_path: str, start: int, end: Optional[int]):
         raise Exception("Requested frames are outbound")
     start = max(start, 0)
     end = min(end, total_frames) if end is not None else total_frames
+    
+    if manual_seek:
+        while start > 0:
+            success = video.grab()
+            if not success:
+                break
+            start -= 1
+    elif start > 0:
+        video.set(cv2.CAP_PROP_POS_FRAMES, start)
+    
     return video, start, end
 
 
 def get_video_frames_generator(
-    source_path: str, stride: int = 1, start: int = 0, end: Optional[int] = None
+    source_path: str, stride: int = 1, start: int = 0, end: Optional[int] = None, manual_seek: bool = False
 ) -> Generator[np.ndarray, None, None]:
     """
     Get a generator that yields the frames of the video.
@@ -144,6 +154,9 @@ def get_video_frames_generator(
             video should generate frames
         end (Optional[int]): Indicates the ending position at which video
             should stop generating frames. If None, video will be read to the end.
+        manual_seek (bool): If True, the generator will manually seek to the
+            `start` frame by grabbing each frame, which is much slower. This is a
+            workaround for videos that don't open at all when you set the 'start' value.
 
     Returns:
         (Generator[np.ndarray, None, None]): A generator that yields the
@@ -157,13 +170,10 @@ def get_video_frames_generator(
             ...
         ```
     """
-    video, start, end = _validate_and_setup_video(source_path, start, end)
-    frame_position = 0
+    video, start, end = _validate_and_setup_video(source_path, start, end, manual_seek)
+    frame_position = start
     while True:
         success, frame = video.read()
-        if frame_position < start:
-            frame_position += 1
-            continue
         if not success or frame_position >= end:
             break
         yield frame