xiaoxue_admin
/
supervision
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(utils): add prefetch to `get_video_frames_generator` (#2273) 

* feat(utils): add prefetch to get_video_frames_generator
* fix(utils): harden _prefetched_frames_generator threading safety
* test(utils): add prefetch combination and minimum-queue tests
* docs(utils): improve prefetch documentation, validation, and test docstrings
* fix(utils): harden prefetch reader-thread exception handling + docs
* docs(changelog): sync front-matter date_modified
* test(utils): harden and extend prefetch test coverage
* test(utils): cover buffered-frames-before-error and zero-frame prefetch cases

---------

Co-authored-by: Jirka Borovec <6035284+Borda@users.noreply.github.com>
Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
This commit is contained in:
Mahbod 2026-07-21 14:12:08 +02:00 committed by GitHub
parent 25e879ec05
commit b20d6eac46
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 405 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
docs/changelog.md
View File

@ -1,6 +1,6 @@
---
description: "Full version history of the supervision Python library — release notes, breaking changes, new features, and deprecations for every version."
date_modified: 2026-07-17
date_modified: 2026-07-21
---
# Changelog
@ -62,6 +62,7 @@ date_modified: 2026-07-17
- Fixed: dataset IO/export edge cases now avoid mutating caller-owned `Detections` during `DetectionDataset` construction, reject non-integer and out-of-range class ids with a clear `ValueError`, load COCO annotations that omit optional `iscrowd`/`area` fields, expose `DetectionDataset.from_coco(use_iscrowd=...)` without changing the existing positional `show_progress` argument, export mask pixel area to COCO when no stored area is present, ignore folder-structure root clutter and non-image files inside class folders, and accept PIL-readable YOLO images such as RGBA or palette PNGs.
### 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.
- Added a cv2-free PyAV fallback for file-video capture, writing, frame seeking,
metadata, and `process_video(preserve_audio=True)` audio remuxing. OpenCV remains
the primary backend when available; `av>=14.2.0` is now required alongside

110
src/supervision/utils/video.py
View File

@ -197,6 +197,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.
@ -212,10 +213,29 @@ 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.
This works best when the consumer releases the GIL during frame
processing (common for numpy/PyTorch/ONNX C-extension calls). Pure
Python per-frame consumers that hold the GIL (for example, heavy
Python loops, PIL usage, or pandas `apply`) usually see little
speedup.
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.
RuntimeError: If `prefetch` is greater than 0 and the background reader
thread encounters a decode/open error, raised as
`RuntimeError(f"Reader thread raised: {item!r}") from item`. Errors are
drained after buffered frames, so when `prefetch` > 0 the consumer may
yield up to `prefetch` additional good frames before the exception is
raised.
Note:
For live camera streams, use `cv2.VideoCapture` with an integer device
@ -246,8 +266,27 @@ def get_video_frames_generator(
for frame in sv.get_video_frames_generator(source_path="<SOURCE_VIDEO_PATH>"):
...
# Prefetch frames in a background thread to overlap I/O with CPU inference:
for frame in sv.get_video_frames_generator(
source_path="<SOURCE_VIDEO_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
)
@ -268,6 +307,73 @@ 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 Exception as exc:
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,

295
tests/utils/test_video.py
View File

@ -1,4 +1,6 @@
import os
import threading
import time
from pathlib import Path
from queue import Empty, Full
from queue import Queue as StdQueue
@ -518,6 +520,102 @@ 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) -> None:
"""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) -> None:
"""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(RuntimeError) as exc_info:
list(get_video_frames_generator(missing_path, prefetch=4))
assert exc_info.value.__cause__ is not None
def test_get_video_frames_generator_prefetch_early_termination(
dummy_video_path,
) -> None:
"""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_releases_on_early_break(monkeypatch) -> None:
"""
Verify that the capture is released when a consumer breaks out early.
@ -554,6 +652,203 @@ def test_get_video_frames_generator_releases_on_early_break(monkeypatch) -> None
assert fake_capture.released
def test_get_video_frames_generator_prefetch_negative_raises(dummy_video_path) -> None:
"""Negative prefetch raises ValueError when the generator is first consumed.
Scenario: Creating get_video_frames_generator with prefetch=-1 and pulling the
first item; because the function is a generator, the guard fires on the first
`next()`, not at call time.
Expected: A ValueError naming the invalid prefetch value is raised.
"""
generator = get_video_frames_generator(dummy_video_path, prefetch=-1)
with pytest.raises(ValueError, match="prefetch must be >= 0"):
next(generator)
def test_get_video_frames_generator_prefetch_stride_early_termination(
dummy_video_path,
) -> None:
"""Breaking early with stride>1 and prefetch>0 exits cleanly without hanging.
Scenario: Consuming only 2 frames from a 10-frame video with stride=2 and
prefetch=4, then creating a fresh strided prefetch generator on the same file.
Expected: The break exits cleanly (no hang); exactly 2 valid frames are taken, and
a fresh strided prefetch generator still yields all 5 strided frames.
"""
taken = []
for frame in get_video_frames_generator(dummy_video_path, stride=2, prefetch=4):
taken.append(frame)
if len(taken) >= 2:
break
assert len(taken) == 2
assert all(isinstance(f, np.ndarray) and f.shape == (480, 640, 3) for f in taken)
# A fresh strided prefetch generator on the same file must still work fully.
fresh = get_video_frames_generator(dummy_video_path, stride=2, prefetch=4)
assert len(list(fresh)) == 5
def test_get_video_frames_generator_prefetch_consumer_exception_cleans_up_thread(
dummy_video_path,
) -> None:
"""A consumer exception propagates and the prefetch reader thread is cleaned up.
Scenario: Raising inside the for-loop body while iterating a prefetch=4 generator,
so the anonymous generator is closed as the exception unwinds.
Expected: The RuntimeError propagates to the caller, and the reader thread is
joined by the generator's finally (active thread count returns to baseline).
"""
def consume_then_raise() -> None:
"""Consume the prefetch generator and raise from inside the loop body."""
for _frame in get_video_frames_generator(dummy_video_path, prefetch=4):
raise RuntimeError("consumer boom")
baseline_threads = threading.active_count()
with pytest.raises(RuntimeError, match="consumer boom"):
consume_then_raise()
assert threading.active_count() == baseline_threads
def test_get_video_frames_generator_prefetch_reader_outlives_join(monkeypatch) -> None:
"""A reader stuck in a slow read must not make the outer generator hang on close.
Scenario: The background reader blocks in a slow `read()` that outlives the
generator's `finally: thread.join(timeout=2.0)`; the consumer closes the
generator after one frame.
Expected: `close()` returns without hanging or raising — the bounded join gives up
and the daemon reader is left to exit on its own.
"""
class SlowCapture:
"""Fake capture that serves one frame then blocks on the next read."""
def __init__(self) -> None:
self.released = False
self.calls = 0
self.block = threading.Event()
def read(self):
"""Return one frame, then block on subsequent reads to simulate a hang."""
self.calls += 1
if self.calls == 1:
return True, np.zeros((2, 2, 3), dtype=np.uint8)
self.block.wait(timeout=10.0)
return False, None
def grab(self):
"""Report a successful grab so stride handling proceeds."""
return True
def release(self) -> None:
"""Record that the capture was released."""
self.released = True
slow_capture = SlowCapture()
monkeypatch.setattr(
"supervision.utils.video._validate_and_setup_video",
lambda *args, **kwargs: (slow_capture, 0, 100),
)
generator = get_video_frames_generator("dummy", prefetch=4)
first_frame = next(generator)
start = time.monotonic()
generator.close()
elapsed = time.monotonic() - start
# Release the daemon reader so it can exit cleanly after the test.
slow_capture.block.set()
assert isinstance(first_frame, np.ndarray)
assert elapsed < 5.0
def test_get_video_frames_generator_prefetch_yields_buffered_frames_before_error(
monkeypatch,
) -> None:
"""Frames already decoded before a mid-stream failure are yielded, then raised.
Scenario: A fake capture successfully reads 3 frames, then raises on the 4th
`read()` call, with `prefetch=4` so all 3 good frames fit in the queue ahead
of the sentinel.
Expected: The consumer receives exactly the 3 good frames in order, then the
wrapping `RuntimeError` — matching the documented "buffered frames before
RuntimeError" ordering guarantee.
"""
class FailAfterNCapture:
"""Fake capture that serves N frames then raises on the next read."""
def __init__(self, good_reads: int) -> None:
self.good_reads = good_reads
self.calls = 0
self.released = False
def read(self):
"""Return a frame for the first `good_reads` calls, then raise."""
self.calls += 1
if self.calls <= self.good_reads:
return True, np.full((2, 2, 3), self.calls, dtype=np.uint8)
raise OSError("simulated mid-stream decode failure")
def grab(self):
"""Report a successful grab so stride handling proceeds."""
return True
def release(self) -> None:
"""Record that the capture was released."""
self.released = True
fake_capture = FailAfterNCapture(good_reads=3)
monkeypatch.setattr(
"supervision.utils.video._validate_and_setup_video",
lambda *args, **kwargs: (fake_capture, 0, 100),
)
collected = []
def consume() -> None:
"""Drain the generator into `collected`, letting the error propagate."""
for frame in get_video_frames_generator("dummy", prefetch=4):
collected.append(frame)
with pytest.raises(RuntimeError) as exc_info:
consume()
assert len(collected) == 3
assert isinstance(exc_info.value.__cause__, OSError)
def test_get_video_frames_generator_prefetch_zero_frame_video(monkeypatch) -> None:
"""The prefetch path yields nothing and returns cleanly for a zero-frame video.
Scenario: A fake capture reports failure on the very first `read()`, simulating
an empty video, with `prefetch=4`.
Expected: The generator yields no frames and returns without hanging on the
initial `frame_queue.get`/`thread.is_alive()` poll.
"""
class EmptyCapture:
"""Fake capture that yields zero frames."""
def read(self):
"""Report immediate end-of-stream."""
return False, None
def grab(self):
"""Report a successful grab so stride handling proceeds."""
return True
def release(self) -> None:
"""No-op release for the empty-video fake."""
monkeypatch.setattr(
"supervision.utils.video._validate_and_setup_video",
lambda *args, **kwargs: (EmptyCapture(), 0, 100),
)
frames = list(get_video_frames_generator("dummy", prefetch=4))
assert frames == []
def test_get_video_frames_generator_with_stride(dummy_video_path) -> None:
"""
Verify that get_video_frames_generator correctly handles the stride parameter.