From 3b09a6356ab030b6bdca5e1aa5db01eb49f21e75 Mon Sep 17 00:00:00 2001 From: realh4m <144081931+realh4m@users.noreply.github.com> Date: Fri, 30 Jan 2026 18:46:19 +0900 Subject: [PATCH] fix: `process_video` silent hang on callback error (#2022) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Handle process_video worker exceptions * test: add unit tests for process_video exception handling and fix docstring typo * fix(pre_commit): 🎨 auto format pre-commit hooks * Handle `Full` exception in video writer queue and add test for small buffer scenario * Apply suggestions from code review --------- Co-authored-by: Your Name Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Jirka Borovec <6035284+Borda@users.noreply.github.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- supervision/utils/video.py | 48 ++++++++++--- test/utils/test_video.py | 133 +++++++++++++++++++++++++++++++++++++ 2 files changed, 172 insertions(+), 9 deletions(-) create mode 100644 test/utils/test_video.py diff --git a/supervision/utils/video.py b/supervision/utils/video.py index f7b8e29c..c4e33da6 100644 --- a/supervision/utils/video.py +++ b/supervision/utils/video.py @@ -5,7 +5,7 @@ import time from collections import deque from collections.abc import Callable, Generator from dataclasses import dataclass -from queue import Queue +from queue import Empty, Full, Queue from typing import Any import cv2 @@ -265,7 +265,7 @@ def process_video( process_video( source_path="source.mp4", target_path="target.mp4", - callback=frame_callback, + callback=callback, ) ``` """ @@ -316,22 +316,52 @@ def process_video( desc=progress_message, ) + exception_in_worker: Exception | None = None + read_finished = False + try: while True: read_item = frame_read_queue.get() if read_item is None: + read_finished = True break frame_index, frame = read_item - processed_frame = callback(frame, frame_index) - - frame_write_queue.put(processed_frame) - progress_bar.update(1) + try: + processed_frame = callback(frame, frame_index) + frame_write_queue.put(processed_frame) + progress_bar.update(1) + except Exception as exc: + exception_in_worker = exc + break finally: - frame_write_queue.put(None) - reader_worker.join() - writer_worker.join() + try: + frame_write_queue.put(None, timeout=1) + except Full: + # Queue is full; this is a best-effort attempt to enqueue the sentinel. + # If we cannot enqueue it, the writer thread will still complete based + # on previously queued frames or other shutdown conditions. + pass + if not read_finished: + while True: + # Use timeout to prevent indefinite blocking if reader thread fails + try: + read_item = frame_read_queue.get(timeout=1) + if read_item is None: + break + # If we timeout waiting for a frame, only assume failure if reader + # thread is no longer alive. Otherwise, keep waiting as the reader + # may simply be slow (for example, due to a slow source). + except Empty: + if not reader_worker.is_alive(): + break + # Reader is still alive; continue waiting for frames. + continue + reader_worker.join(timeout=10) + writer_worker.join(timeout=10) progress_bar.close() + if exception_in_worker is not None: + raise exception_in_worker class FPSMonitor: diff --git a/test/utils/test_video.py b/test/utils/test_video.py new file mode 100644 index 00000000..7c9e1d98 --- /dev/null +++ b/test/utils/test_video.py @@ -0,0 +1,133 @@ +import os + +import cv2 +import numpy as np +import pytest + +from supervision.utils.video import VideoInfo, get_video_frames_generator, process_video + + +@pytest.fixture +def dummy_video_path(tmp_path): + path = str(tmp_path / "dummy_video.mp4") + fourcc = cv2.VideoWriter_fourcc(*"mp4v") + out = cv2.VideoWriter(path, fourcc, 25, (640, 480)) + for _ in range(10): + frame = np.zeros((480, 640, 3), dtype=np.uint8) + out.write(frame) + out.release() + return path + + +def test_process_video_exception_handling(dummy_video_path, tmp_path): + target_path = str(tmp_path / "target.mp4") + + def callback_with_exception(frame, index): + if index == 5: + raise ValueError("Test exception at frame 5") + return frame + + with pytest.raises(ValueError, match="Test exception at frame 5"): + process_video( + source_path=dummy_video_path, + target_path=target_path, + callback=callback_with_exception, + ) + + +def test_process_video_success(dummy_video_path, tmp_path): + target_path = str(tmp_path / "target_success.mp4") + + def callback_success(frame, index): + return frame + + # This should complete without exception + process_video( + source_path=dummy_video_path, target_path=target_path, callback=callback_success + ) + + assert os.path.exists(target_path) + + +def test_process_video_exception_with_small_buffer(dummy_video_path, tmp_path): + target_path = str(tmp_path / "target_exception_small_buffer.mp4") + + def callback_with_exception(frame, index): + if index == 5: + raise ValueError("Test exception at frame 5") + return frame + + with pytest.raises(ValueError, match="Test exception at frame 5"): + process_video( + source_path=dummy_video_path, + target_path=target_path, + callback=callback_with_exception, + prefetch=1, + writer_buffer=1, + ) + + +def test_process_video_max_frames(dummy_video_path, tmp_path): + target_path = str(tmp_path / "target_max_frames.mp4") + processed_indices = [] + + def callback(frame, index): + processed_indices.append(index) + return frame + + process_video( + source_path=dummy_video_path, + target_path=target_path, + callback=callback, + max_frames=5, + ) + + assert len(processed_indices) == 5 + assert processed_indices == [0, 1, 2, 3, 4] + + +def test_process_video_custom_params(dummy_video_path, tmp_path): + target_path = str(tmp_path / "target_custom_params.mp4") + + def callback(frame, index): + return frame + + # Test with very small prefetch and writer_buffer + process_video( + source_path=dummy_video_path, + target_path=target_path, + callback=callback, + prefetch=1, + writer_buffer=1, + ) + + assert os.path.exists(target_path) + + +def test_video_info(dummy_video_path): + video_info = VideoInfo.from_video_path(dummy_video_path) + assert video_info.width == 640 + assert video_info.height == 480 + assert video_info.fps == 25 + assert video_info.total_frames == 10 + assert video_info.resolution_wh == (640, 480) + + +def test_get_video_frames_generator(dummy_video_path): + generator = get_video_frames_generator(dummy_video_path) + frames = list(generator) + assert len(frames) == 10 + assert all(isinstance(frame, np.ndarray) for frame in frames) + assert all(frame.shape == (480, 640, 3) for frame in frames) + + +def test_get_video_frames_generator_with_stride(dummy_video_path): + generator = get_video_frames_generator(dummy_video_path, stride=2) + frames = list(generator) + assert len(frames) == 5 + + +def test_get_video_frames_generator_with_start_end(dummy_video_path): + generator = get_video_frames_generator(dummy_video_path, start=2, end=5) + frames = list(generator) + assert len(frames) == 3