211 lines
6.7 KiB
Python
211 lines
6.7 KiB
Python
"""Tests for the PyAV-backed video compatibility surface."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import av
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from supervision import _cv2
|
|
from supervision._cv2._video import (
|
|
_mux_audio,
|
|
_VideoCapture,
|
|
_VideoWriter,
|
|
)
|
|
|
|
|
|
def _write_video(path: Path, values: list[int], fps: int = 5) -> None:
|
|
"""Write a small deterministic MPEG-4 video for fallback tests."""
|
|
container = av.open(str(path), mode="w")
|
|
stream = container.add_stream("mpeg4", rate=fps)
|
|
stream.width = 16
|
|
stream.height = 16
|
|
stream.pix_fmt = "yuv420p"
|
|
try:
|
|
for value in values:
|
|
frame = av.VideoFrame.from_ndarray(
|
|
np.full((16, 16, 3), value, dtype=np.uint8), format="bgr24"
|
|
)
|
|
for packet in stream.encode(frame):
|
|
container.mux(packet)
|
|
for packet in stream.encode():
|
|
container.mux(packet)
|
|
finally:
|
|
container.close()
|
|
|
|
|
|
def _write_video_with_audio(path: Path, frame_count: int = 5, fps: int = 5) -> None:
|
|
"""Write a short video with one AAC audio stream for remux tests."""
|
|
container = av.open(str(path), mode="w")
|
|
video_stream = container.add_stream("mpeg4", rate=fps)
|
|
video_stream.width = 16
|
|
video_stream.height = 16
|
|
video_stream.pix_fmt = "yuv420p"
|
|
audio_stream = container.add_stream("aac", rate=8_000)
|
|
audio_stream.layout = "mono"
|
|
try:
|
|
for value in range(frame_count):
|
|
frame = av.VideoFrame.from_ndarray(
|
|
np.full((16, 16, 3), value * 20, dtype=np.uint8), format="bgr24"
|
|
)
|
|
for packet in video_stream.encode(frame):
|
|
container.mux(packet)
|
|
|
|
samples = np.zeros((1, 8_000), dtype=np.int16)
|
|
audio_frame = av.AudioFrame.from_ndarray(samples, format="s16", layout="mono")
|
|
audio_frame.sample_rate = 8_000
|
|
audio_frame.pts = 0
|
|
for packet in audio_stream.encode(audio_frame):
|
|
container.mux(packet)
|
|
for packet in video_stream.encode():
|
|
container.mux(packet)
|
|
for packet in audio_stream.encode():
|
|
container.mux(packet)
|
|
finally:
|
|
container.close()
|
|
|
|
|
|
def _run_without_opencv(source: str) -> None:
|
|
"""Run a Python snippet with cv2 imports blocked."""
|
|
env = os.environ.copy()
|
|
source_path = str(Path(__file__).resolve().parents[2] / "src")
|
|
env["PYTHONPATH"] = os.pathsep.join(
|
|
filter(None, (source_path, env.get("PYTHONPATH")))
|
|
)
|
|
result = subprocess.run(
|
|
[sys.executable, "-c", source],
|
|
check=False,
|
|
capture_output=True,
|
|
text=True,
|
|
env=env,
|
|
timeout=60,
|
|
)
|
|
assert result.returncode == 0, result.stderr
|
|
|
|
|
|
def test_video_module_uses_required_pyav_dependency() -> None:
|
|
"""The video fallback imports required PyAV directly without a loader."""
|
|
from supervision._cv2 import _video
|
|
|
|
assert _video.av.__name__ == "av"
|
|
assert not hasattr(_video, "_load_av")
|
|
|
|
|
|
def test_fallback_capture_reports_metadata_and_supports_exact_seek(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""Fallback capture exposes metadata and starts decoding at requested frames."""
|
|
source_path = tmp_path / "source.mp4"
|
|
_write_video(source_path, [0, 40, 80, 120, 160])
|
|
|
|
capture = _VideoCapture(str(source_path))
|
|
assert capture.isOpened()
|
|
assert capture.get(_cv2.CAP_PROP_FRAME_WIDTH) == 16
|
|
assert capture.get(_cv2.CAP_PROP_FRAME_HEIGHT) == 16
|
|
assert capture.get(_cv2.CAP_PROP_FPS) == pytest.approx(5.0)
|
|
assert capture.get(_cv2.CAP_PROP_FRAME_COUNT) == 5
|
|
|
|
assert capture.set(_cv2.CAP_PROP_POS_FRAMES, 3)
|
|
success, frame = capture.read()
|
|
capture.release()
|
|
|
|
assert success
|
|
assert frame is not None
|
|
assert frame.shape == (16, 16, 3)
|
|
assert float(frame.mean()) == pytest.approx(120.0, abs=20.0)
|
|
assert not capture.isOpened()
|
|
|
|
|
|
def test_fallback_writer_default_codec_round_trips(tmp_path: Path) -> None:
|
|
"""The guaranteed mp4v fallback writer creates a readable video."""
|
|
target_path = tmp_path / "target.mp4"
|
|
fourcc = _cv2.VideoWriter_fourcc(*"mp4v")
|
|
writer = _VideoWriter(str(target_path), fourcc, 5.0, (16, 16))
|
|
|
|
assert writer.isOpened()
|
|
for value in [0, 40, 80]:
|
|
writer.write(np.full((16, 16, 3), value, dtype=np.uint8))
|
|
writer.release()
|
|
|
|
capture = _VideoCapture(str(target_path))
|
|
frames = []
|
|
while True:
|
|
success, frame = capture.read()
|
|
if not success:
|
|
break
|
|
assert frame is not None
|
|
frames.append(frame)
|
|
capture.release()
|
|
|
|
assert len(frames) == 3
|
|
assert target_path.stat().st_size > 0
|
|
|
|
|
|
def test_fallback_video_works_when_opencv_is_blocked(tmp_path: Path) -> None:
|
|
"""Production video APIs use PyAV when cv2 cannot be imported."""
|
|
source_path = tmp_path / "source.mp4"
|
|
target_path = tmp_path / "target.mp4"
|
|
_write_video(source_path, [0, 40, 80])
|
|
source = f"""
|
|
import sys
|
|
|
|
|
|
class BlockCv2:
|
|
def find_spec(self, fullname, path=None, target=None):
|
|
if fullname == "cv2":
|
|
raise ModuleNotFoundError("blocked for test")
|
|
return None
|
|
|
|
|
|
sys.meta_path.insert(0, BlockCv2())
|
|
from supervision import _cv2
|
|
from supervision.utils.video import VideoInfo, VideoSink, get_video_frames_generator
|
|
|
|
assert _cv2._IS_CV2_AVAILABLE is False
|
|
info = VideoInfo.from_video_path({str(source_path)!r})
|
|
assert (info.width, info.height, info.total_frames) == (16, 16, 3)
|
|
frames = list(get_video_frames_generator({str(source_path)!r}, start=1, end=3))
|
|
assert len(frames) == 2
|
|
with VideoSink({str(target_path)!r}, info) as sink:
|
|
for frame in frames:
|
|
sink.write_frame(frame)
|
|
assert _cv2.VideoCapture({str(target_path)!r}).isOpened()
|
|
"""
|
|
_run_without_opencv(source)
|
|
|
|
|
|
def test_mux_audio_remuxes_first_audio_stream_and_truncates_to_video(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""Audio remuxing uses PyAV and keeps the processed video duration."""
|
|
source_path = tmp_path / "source_with_audio.mp4"
|
|
target_path = tmp_path / "target.mp4"
|
|
_write_video_with_audio(source_path, frame_count=5)
|
|
_write_video(target_path, [0, 40, 80], fps=5)
|
|
|
|
_mux_audio(str(source_path), str(target_path))
|
|
|
|
output = av.open(str(target_path))
|
|
try:
|
|
assert len(output.streams.video) == 1
|
|
assert len(output.streams.audio) == 1
|
|
assert output.streams.video[0].frames == 3
|
|
finally:
|
|
output.close()
|
|
|
|
|
|
def test_mux_audio_leaves_target_unchanged_on_failure(tmp_path: Path) -> None:
|
|
"""A failed PyAV remux never replaces the existing target file."""
|
|
target_path = tmp_path / "target.mp4"
|
|
_write_video(target_path, [0, 40, 80])
|
|
original = target_path.read_bytes()
|
|
|
|
_mux_audio(str(tmp_path / "missing.mp4"), str(target_path))
|
|
|
|
assert target_path.read_bytes() == original
|