342 lines
13 KiB
Python
342 lines
13 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""Tests for omni_duplex_eval_media — PyAV-backed media helpers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
from pathlib import Path
|
|
|
|
import av
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from vllm_omni.benchmarks.duplex.omni_duplex_eval_media import (
|
|
_ffmpeg_q_to_pil_quality,
|
|
extract_jpeg,
|
|
materialize_media,
|
|
read_audio_pcm16,
|
|
video_duration,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers: synthetic test media generated with PyAV (no external files)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_synthetic_mp4(duration: float = 1.0, fps: int = 30) -> bytes:
|
|
"""Return an MP4 byte blob with a solid-color pattern.
|
|
|
|
Uses the built-in ``mpeg4`` encoder — available everywhere PyAV's bundled
|
|
FFmpeg is present — so no external codec dependency.
|
|
"""
|
|
buf = io.BytesIO()
|
|
total_frames = int(duration * fps)
|
|
with av.open(buf, mode="w", format="mp4") as container:
|
|
stream = container.add_stream("mpeg4", rate=fps)
|
|
stream.width = 64
|
|
stream.height = 64
|
|
stream.pix_fmt = "yuv420p"
|
|
for i in range(total_frames):
|
|
frame = av.VideoFrame.from_ndarray(np.full((64, 64, 3), (i * 8) % 256, dtype=np.uint8), format="rgb24")
|
|
for pkt in stream.encode(frame):
|
|
container.mux(pkt)
|
|
for pkt in stream.encode():
|
|
container.mux(pkt)
|
|
return buf.getvalue()
|
|
|
|
|
|
def _make_synthetic_wav(duration: float = 1.0, sample_rate: int = 16000) -> bytes:
|
|
"""Return a 16-bit mono PCM WAV byte blob (440 Hz sine wave)."""
|
|
import wave as _wave
|
|
|
|
n = int(sample_rate * duration)
|
|
samples = (np.sin(2 * np.pi * 440 * np.arange(n) / sample_rate) * 32767).astype(np.int16)
|
|
buf = io.BytesIO()
|
|
with _wave.open(buf, "wb") as wf:
|
|
wf.setnchannels(1)
|
|
wf.setsampwidth(2)
|
|
wf.setframerate(sample_rate)
|
|
wf.writeframes(samples.tobytes())
|
|
return buf.getvalue()
|
|
|
|
|
|
def _make_synthetic_ogg(duration: float = 1.0, sample_rate: int = 16000) -> bytes:
|
|
"""Return an OGG Vorbis byte blob (440 Hz sine wave).
|
|
|
|
Uses ``libopus`` as a fallback if ``libvorbis`` is unavailable in the
|
|
bundled FFmpeg (e.g. PyAV wheels from an older build).
|
|
"""
|
|
buf = io.BytesIO()
|
|
n = int(sample_rate * duration)
|
|
samples = (np.sin(2 * np.pi * 440 * np.arange(n) / sample_rate) * 0.3).astype(np.float32)
|
|
with av.open(buf, mode="w", format="ogg") as container:
|
|
codec_name = "libopus" if "libvorbis" not in av.codecs_available else "libvorbis"
|
|
stream = container.add_stream(codec_name, rate=sample_rate, layout="mono")
|
|
frame = av.AudioFrame.from_ndarray(samples[np.newaxis, :], format="fltp", layout="mono")
|
|
frame.sample_rate = sample_rate
|
|
for pkt in stream.encode(frame):
|
|
container.mux(pkt)
|
|
for pkt in stream.encode():
|
|
container.mux(pkt)
|
|
return buf.getvalue()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# video_duration
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestVideoDuration:
|
|
@staticmethod
|
|
def test_returns_duration(tmp_path: Path) -> None:
|
|
mp4 = _make_synthetic_mp4(duration=2.0, fps=30)
|
|
path = tmp_path / "test.mp4"
|
|
path.write_bytes(mp4)
|
|
dur = video_duration(path)
|
|
assert abs(dur - 2.0) < 0.05, f"expected ~2.0s, got {dur}"
|
|
|
|
@staticmethod
|
|
def test_raises_for_missing_file(tmp_path: Path) -> None:
|
|
with pytest.raises((OSError, ValueError)):
|
|
video_duration(tmp_path / "nonexistent.mp4")
|
|
|
|
@staticmethod
|
|
def test_raises_for_broken_file(tmp_path: Path) -> None:
|
|
broken = tmp_path / "broken.mp4"
|
|
broken.write_bytes(b"\x00\x00\x00\x00")
|
|
with pytest.raises((OSError, ValueError)):
|
|
video_duration(broken)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# extract_jpeg
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestExtractJpeg:
|
|
@staticmethod
|
|
def test_returns_valid_jpeg(tmp_path: Path) -> None:
|
|
mp4 = _make_synthetic_mp4(duration=1.0, fps=30)
|
|
path = tmp_path / "test.mp4"
|
|
path.write_bytes(mp4)
|
|
jpeg = extract_jpeg(path, timestamp=0.5, quality=3)
|
|
assert jpeg.startswith(b"\xff\xd8"), "not a valid JPEG header"
|
|
assert len(jpeg) > 100
|
|
|
|
@staticmethod
|
|
def test_quality_maps_to_pil_scale(tmp_path: Path) -> None:
|
|
"""Higher ffmpeg quality (lower q value) produces a larger JPEG."""
|
|
mp4 = _make_synthetic_mp4(duration=1.0, fps=30)
|
|
path = tmp_path / "test.mp4"
|
|
path.write_bytes(mp4)
|
|
high = extract_jpeg(path, timestamp=0.5, quality=1)
|
|
low = extract_jpeg(path, timestamp=0.5, quality=31)
|
|
assert len(high) > len(low), "quality=1 should produce larger JPEG than quality=31"
|
|
|
|
@staticmethod
|
|
def test_negative_timestamp_clamped(tmp_path: Path) -> None:
|
|
"""Negative timestamps are clamped to 0."""
|
|
mp4 = _make_synthetic_mp4(duration=1.0, fps=30)
|
|
path = tmp_path / "test.mp4"
|
|
path.write_bytes(mp4)
|
|
jpeg = extract_jpeg(path, timestamp=-1.0, quality=3)
|
|
assert jpeg.startswith(b"\xff\xd8")
|
|
|
|
@staticmethod
|
|
def test_raises_value_error_for_no_video(tmp_path: Path) -> None:
|
|
wav = _make_synthetic_wav()
|
|
path = tmp_path / "test.wav"
|
|
path.write_bytes(wav)
|
|
with pytest.raises(ValueError):
|
|
extract_jpeg(path, timestamp=0.0)
|
|
|
|
@staticmethod
|
|
def test_long_gop_returns_different_frames_for_different_timestamps(tmp_path: Path) -> None:
|
|
"""Long GOP: different timestamps must return different frame content (P2)."""
|
|
# Create a long-GOP video: 3s @ 10fps = 30 frames, gop_size=30 (1 keyframe)
|
|
buf = io.BytesIO()
|
|
fps = 10
|
|
duration = 3.0
|
|
total_frames = int(duration * fps)
|
|
with av.open(buf, mode="w", format="mp4") as container:
|
|
stream = container.add_stream("mpeg4", rate=fps)
|
|
stream.width = 64
|
|
stream.height = 64
|
|
stream.pix_fmt = "yuv420p"
|
|
stream.gop_size = 30 # 3 seconds per keyframe
|
|
for i in range(total_frames):
|
|
frame = av.VideoFrame.from_ndarray(np.full((64, 64, 3), (i * 64) % 256, dtype=np.uint8), format="rgb24")
|
|
for pkt in stream.encode(frame):
|
|
container.mux(pkt)
|
|
for pkt in stream.encode():
|
|
container.mux(pkt)
|
|
mp4 = buf.getvalue()
|
|
|
|
path = tmp_path / "long_gop.mp4"
|
|
path.write_bytes(mp4)
|
|
|
|
# Extract frames at different timestamps
|
|
frame0 = extract_jpeg(path, timestamp=0.0, quality=3)
|
|
frame1 = extract_jpeg(path, timestamp=1.0, quality=3)
|
|
frame2 = extract_jpeg(path, timestamp=2.0, quality=3)
|
|
|
|
# Verify all are valid JPEGs
|
|
assert frame0.startswith(b"\xff\xd8")
|
|
assert frame1.startswith(b"\xff\xd8")
|
|
assert frame2.startswith(b"\xff\xd8")
|
|
|
|
# Verify content differs: decode JPEGs and compare mean brightness
|
|
from PIL import Image
|
|
|
|
img0 = np.asarray(Image.open(io.BytesIO(frame0)))
|
|
img1 = np.asarray(Image.open(io.BytesIO(frame1)))
|
|
img2 = np.asarray(Image.open(io.BytesIO(frame2)))
|
|
means = [img.mean() for img in (img0, img1, img2)]
|
|
assert len(set(round(m, 1) for m in means)) >= 2, (
|
|
f"timestamps 0/1/2s should give different frames, got means {means}"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# read_audio_pcm16
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestReadAudioPcm16:
|
|
@staticmethod
|
|
def test_wav_fast_path(tmp_path: Path) -> None:
|
|
"""WAV with matching params uses the fast wave module path."""
|
|
wav = _make_synthetic_wav(duration=1.0, sample_rate=16000)
|
|
path = tmp_path / "test.wav"
|
|
path.write_bytes(wav)
|
|
pcm = read_audio_pcm16(path)
|
|
assert len(pcm) == 16000 * 2, f"expected 32000 bytes, got {len(pcm)}"
|
|
assert isinstance(pcm, bytes)
|
|
|
|
@staticmethod
|
|
def test_wav_different_rate_resamples(tmp_path: Path) -> None:
|
|
"""A 48 kHz WAV is resampled down to 16 kHz — exact sample count after flush."""
|
|
wav = _make_synthetic_wav(duration=1.0, sample_rate=48000)
|
|
path = tmp_path / "test.wav"
|
|
path.write_bytes(wav)
|
|
pcm = read_audio_pcm16(path)
|
|
assert len(pcm) > 0
|
|
assert len(pcm) % 2 == 0
|
|
# 1s @ 48kHz → 1s @ 16kHz = 16000 samples = 32000 bytes (±1 sample)
|
|
assert abs(len(pcm) - 32000) <= 4, f"expected ~32000 bytes, got {len(pcm)}"
|
|
|
|
@staticmethod
|
|
def test_non_wav_ogg(tmp_path: Path) -> None:
|
|
"""Non-WAV audio (OGG) is decoded and resampled via PyAV."""
|
|
ogg = _make_synthetic_ogg(duration=1.0, sample_rate=16000)
|
|
path = tmp_path / "test.ogg"
|
|
path.write_bytes(ogg)
|
|
pcm = read_audio_pcm16(path)
|
|
assert len(pcm) > 0
|
|
assert len(pcm) % 2 == 0
|
|
|
|
@staticmethod
|
|
def test_non_wav_ogg_resample_accuracy(tmp_path: Path) -> None:
|
|
"""OGG resample should produce exact sample count after flush."""
|
|
ogg = _make_synthetic_ogg(duration=1.5, sample_rate=48000)
|
|
path = tmp_path / "test.ogg"
|
|
path.write_bytes(ogg)
|
|
pcm = read_audio_pcm16(path)
|
|
# 1.5s @ 48kHz → 1.5s @ 16kHz = 24000 samples = 48000 bytes (±1 sample)
|
|
expected = int(1.5 * 16000) * 2
|
|
assert abs(len(pcm) - expected) <= 4, f"expected ~{expected} bytes, got {len(pcm)}"
|
|
|
|
@staticmethod
|
|
def test_raises_for_broken_file(tmp_path: Path) -> None:
|
|
"""Corrupt audio must raise, not return empty bytes (L1)."""
|
|
broken = tmp_path / "broken.mp3"
|
|
broken.write_bytes(b"\x00\x00\x00\x00")
|
|
with pytest.raises((OSError, ValueError)):
|
|
read_audio_pcm16(broken)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _ffmpeg_q_to_pil_quality
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestFfmpegQToPilQuality:
|
|
@staticmethod
|
|
def test_q1_maps_to_best() -> None:
|
|
assert _ffmpeg_q_to_pil_quality(1) >= 90
|
|
|
|
@staticmethod
|
|
def test_q31_maps_to_lowest() -> None:
|
|
assert _ffmpeg_q_to_pil_quality(31) <= 10
|
|
|
|
@staticmethod
|
|
def test_monotonic() -> None:
|
|
"""Higher ffmpeg q → lower PIL quality (monotonically non-increasing)."""
|
|
prev = 100
|
|
for q in range(1, 32):
|
|
p = _ffmpeg_q_to_pil_quality(q)
|
|
assert p <= prev, f"quality.increase at q={q}: {prev} → {p}"
|
|
prev = p
|
|
|
|
|
|
def test_dict_with_bytes_materializes_a_real_file(tmp_path):
|
|
# A Hugging Face Audio/Video feature value {"bytes": ..., "path": ...} used
|
|
# to short-circuit on ``path`` (the artifact name inside the dataset, which
|
|
# does not exist on disk), so downstream ffmpeg saw a missing file. Inline
|
|
# bytes must win and be written to a real file whose content equals bytes.
|
|
payload = b"\x00\xff\x01RIFFfake-content"
|
|
resolved = materialize_media({"bytes": payload, "path": "question_audio.wav"}, tmp_path, "s1_question", ".wav")
|
|
assert resolved.is_file()
|
|
assert resolved.read_bytes() == payload
|
|
assert resolved.parent == tmp_path
|
|
assert resolved.name == "s1_question.wav"
|
|
|
|
|
|
def test_dict_extension_comes_from_path_safely(tmp_path):
|
|
payload = b"flac-ish-bytes"
|
|
resolved = materialize_media({"bytes": payload, "path": "audio/clip.FLAC"}, tmp_path, "s1", ".wav")
|
|
assert resolved.name == "s1.flac"
|
|
assert resolved.read_bytes() == payload
|
|
# A path with no usable extension falls back to the caller's suffix.
|
|
resolved = materialize_media({"bytes": payload, "path": "audio/clip"}, tmp_path, "s2", ".wav")
|
|
assert resolved.name == "s2.wav"
|
|
assert resolved.read_bytes() == payload
|
|
|
|
|
|
def test_materialization_is_idempotent_and_updated(tmp_path):
|
|
payload = b"abc"
|
|
first = materialize_media({"bytes": payload, "path": "q.wav"}, tmp_path, "s", ".wav")
|
|
second = materialize_media({"bytes": payload, "path": "q.wav"}, tmp_path, "s", ".wav")
|
|
assert first == second
|
|
assert first.read_bytes() == payload
|
|
replacement = b"xyz"
|
|
third = materialize_media({"bytes": replacement, "path": "q.wav"}, tmp_path, "s", ".wav")
|
|
assert third == first
|
|
assert third.read_bytes() == replacement
|
|
|
|
|
|
def test_str_path_and_path_only_mapping_returned_as_is(tmp_path):
|
|
existing = tmp_path / "clip.mp4"
|
|
existing.write_bytes(b"video")
|
|
assert materialize_media(str(existing), tmp_path, "s", ".mp4") == existing
|
|
assert materialize_media({"path": str(existing)}, tmp_path, "s", ".mp4") == existing
|
|
|
|
|
|
def test_raw_bytes_value_is_written(tmp_path):
|
|
payload = b"raw-media"
|
|
resolved = materialize_media(payload, tmp_path, "r", ".bin")
|
|
assert resolved.name == "r.bin"
|
|
assert resolved.read_bytes() == payload
|
|
|
|
|
|
def test_unresolvable_value_raises_clear_error(tmp_path):
|
|
with pytest.raises(ValueError, match="cannot materialize"):
|
|
materialize_media(None, tmp_path, "s", ".wav")
|
|
with pytest.raises(ValueError, match="cannot materialize"):
|
|
materialize_media({"path": None}, tmp_path, "s", ".wav")
|