diff --git a/docs/changelog.md b/docs/changelog.md index 0a0f2552f..e0149715e 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -17,6 +17,7 @@ date_modified: 2026-06-25 - Fixed [#2382](https://github.com/roboflow/supervision/pull/2382): `sv.Detections.get_anchors_coordinates` now uses oriented bounding box corners (`data["xyxyxyxy"]`) when OBB data is present, instead of falling back to the axis-aligned envelope. Anchors on rotated detections now lie on the oriented body rather than drifting to the envelope. Non-OBB detections and `Position.CENTER_OF_MASS` (which requires a mask) are unaffected. ### Added +- `sv.get_video_frames_generator` now accepts `prefetch: int = 0` ([#2273](https://github.com/roboflow/supervision/pull/2273)). When `> 0`, frames are decoded on a background daemon thread and buffered in a bounded queue, overlapping I/O with consumer processing. Default `0` preserves the existing synchronous behaviour. - `BaseAnnotator.requires_mask` — class-level `bool` flag on all annotators; `True` for `MaskAnnotator`, `PolygonAnnotator`, and `HaloAnnotator`; `False` for all others. Integrations can inspect this before materializing expensive mask payloads ([#2370](https://github.com/roboflow/supervision/pull/2370)) - `CompactMask.from_coco_rle` — efficient COCO RLE ingestion into crop-scoped compact mask format without materializing dense `(N, H, W)` arrays ([#2367](https://github.com/roboflow/supervision/pull/2367)) - `Detections.from_inference(compact_masks=True)` — opt-in compact mask representation for Roboflow/Inference segmentation results; masks are cropped to detector bounding boxes ([#2367](https://github.com/roboflow/supervision/pull/2367)) diff --git a/src/supervision/utils/video.py b/src/supervision/utils/video.py index 51a2e16d2..a4f336b04 100644 --- a/src/supervision/utils/video.py +++ b/src/supervision/utils/video.py @@ -241,6 +241,7 @@ def get_video_frames_generator( start: int = 0, end: int | None = None, iterative_seek: bool = False, + prefetch: int = 0, ) -> Generator[npt.NDArray[np.uint8], None, None]: """ Get a generator that yields the frames of the video. @@ -256,10 +257,18 @@ def get_video_frames_generator( iterative_seek: If True, the generator will 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. + prefetch: If > 0, decode frames in a background thread and buffer up to + this many frames in a bounded queue. Useful when the consumer (e.g. + CPU inference) is the bottleneck and can overlap with decode I/O. + Default 0 keeps the original synchronous behaviour unchanged. Note: + each buffered frame occupies width x height x 3 bytes of uncompressed + memory; use `sv.VideoInfo.from_video_path()` to size appropriately. Returns: - A generator that yields the - frames of the video. + A generator that yields the frames of the video. + + Raises: + ValueError: If `prefetch` is negative. Examples: ```python @@ -267,8 +276,27 @@ def get_video_frames_generator( for frame in sv.get_video_frames_generator(source_path=""): ... + + # Prefetch frames in a background thread to overlap I/O with CPU inference: + for frame in sv.get_video_frames_generator( + source_path="", prefetch=8 + ): + ... ``` """ + if prefetch < 0: + raise ValueError(f"prefetch must be >= 0, got {prefetch!r}") + if prefetch > 0: + yield from _prefetched_frames_generator( + source_path=source_path, + stride=stride, + start=start, + end=end, + iterative_seek=iterative_seek, + prefetch=prefetch, + ) + return + video, start, end = _validate_and_setup_video( source_path, start, end, iterative_seek ) @@ -287,6 +315,74 @@ def get_video_frames_generator( video.release() +def _prefetched_frames_generator( + source_path: str, + stride: int, + start: int, + end: int | None, + iterative_seek: bool, + prefetch: int, +) -> Generator[npt.NDArray[np.uint8], None, None]: + """Read frames into a bounded queue on a daemon thread. + + Sentinel protocol: None = normal EOF, Exception instance = reader error. + """ + frame_queue: Queue[npt.NDArray[np.uint8] | BaseException | None] = Queue( + maxsize=prefetch + ) + stop_event = threading.Event() + + def reader() -> None: + sentinel: BaseException | None = None + try: + for frame in get_video_frames_generator( + source_path=source_path, + stride=stride, + start=start, + end=end, + iterative_seek=iterative_seek, + prefetch=0, + ): + if stop_event.is_set(): + return + while True: + try: + frame_queue.put(frame, timeout=0.1) + break + except Full: + if stop_event.is_set(): + return + except BaseException as exc: + if isinstance(exc, Exception): + sentinel = exc + finally: + while not stop_event.is_set(): + try: + frame_queue.put(sentinel, timeout=0.1) + return + except Full: + pass + + thread = threading.Thread(target=reader, daemon=True) + thread.start() + try: + while True: + try: + item = frame_queue.get(timeout=0.5) + except Empty: + if not thread.is_alive(): + break + continue + if isinstance(item, BaseException): + raise RuntimeError(f"Reader thread raised: {item!r}") from item + if item is None: + break + yield item + finally: + stop_event.set() + thread.join(timeout=2.0) + + def process_video( source_path: str, target_path: str, diff --git a/tests/utils/test_video.py b/tests/utils/test_video.py index 2310ec802..79a0cab78 100644 --- a/tests/utils/test_video.py +++ b/tests/utils/test_video.py @@ -200,6 +200,99 @@ def test_get_video_frames_generator(dummy_video_path) -> None: assert all(frame.shape == (480, 640, 3) for frame in frames) +def test_get_video_frames_generator_prefetch_matches_sync(dummy_video_path): + """Verify that the prefetch path yields identical frames to the sync path. + + Scenario: Iterating over a video with prefetch=4 and again with prefetch=0 + (synchronous) on the same dummy video. + Expected: Both generators yield the same number of frames in the same order, + with each corresponding frame being pixel-for-pixel identical. + """ + sync_frames = list(get_video_frames_generator(dummy_video_path)) + prefetched_frames = list(get_video_frames_generator(dummy_video_path, prefetch=4)) + assert len(prefetched_frames) == len(sync_frames) == 10 + for a, b in zip(prefetched_frames, sync_frames): + assert np.array_equal(a, b) + + +def test_get_video_frames_generator_prefetch_propagates_decode_errors(tmp_path): + """Verify that reader-thread exceptions reach the consumer, not get swallowed. + + Scenario: Passing a non-existent file path to the prefetch path so the reader + thread fails immediately on video open. + Expected: The exception propagates to the consumer and is raised as a + RuntimeError wrapping the original error; the consumer does not hang. + """ + missing_path = str(tmp_path / "does_not_exist.mp4") + with pytest.raises((Exception, RuntimeError)): + list(get_video_frames_generator(missing_path, prefetch=4)) + + +def test_get_video_frames_generator_prefetch_early_termination(dummy_video_path): + """Verify that breaking out of the prefetched generator does not block reuse. + + Scenario: Consuming only 3 frames from a 10-frame video with prefetch=4, then + creating a fresh generator on the same file. + Expected: The break exits cleanly without hanging; a new generator on the same + file yields all 10 frames normally. + """ + taken = [] + for frame in get_video_frames_generator(dummy_video_path, prefetch=4): + taken.append(frame) + if len(taken) >= 3: + break + assert len(taken) == 3 + # A fresh generator on the same file must still work normally. + assert len(list(get_video_frames_generator(dummy_video_path, prefetch=4))) == 10 + + +@pytest.mark.parametrize( + ("stride", "start", "end"), + [ + pytest.param(2, 0, None, id="stride2"), + pytest.param(1, 2, 7, id="start2_end7"), + pytest.param(2, 2, 8, id="stride2_start2_end8"), + ], +) +def test_get_video_frames_generator_prefetch_param_forwarding( + dummy_video_path, stride, start, end +) -> None: + """Prefetch path must forward stride/start/end identically to the sync path. + + Scenario: Using the prefetch path with various stride, start, and end + combinations to verify parameters are correctly forwarded. + Expected: The prefetch output matches the sync path frame-for-frame for + each combination; no frames skipped or duplicated. + """ + sync_frames = list( + get_video_frames_generator( + dummy_video_path, stride=stride, start=start, end=end + ) + ) + prefetched_frames = list( + get_video_frames_generator( + dummy_video_path, stride=stride, start=start, end=end, prefetch=4 + ) + ) + assert len(prefetched_frames) == len(sync_frames) + for a, b in zip(prefetched_frames, sync_frames): + assert np.array_equal(a, b) + + +def test_get_video_frames_generator_prefetch_minimum_queue(dummy_video_path) -> None: + """prefetch=1 creates maximum backpressure; all frames must be returned in order. + + Scenario: Using prefetch=1 forces the reader to block after every decoded + frame, maximising producer-consumer synchronisation pressure. + Expected: All 10 frames are yielded in the same order as the sync path. + """ + sync_frames = list(get_video_frames_generator(dummy_video_path)) + prefetched_frames = list(get_video_frames_generator(dummy_video_path, prefetch=1)) + assert len(prefetched_frames) == len(sync_frames) == 10 + for a, b in zip(prefetched_frames, sync_frames): + assert np.array_equal(a, b) + + def test_get_video_frames_generator_with_stride(dummy_video_path) -> None: """ Verify that get_video_frames_generator correctly handles the stride parameter.