diff --git a/mmif/utils/audio_document_helper.py b/mmif/utils/audio_document_helper.py new file mode 100644 index 00000000..f622ee5c --- /dev/null +++ b/mmif/utils/audio_document_helper.py @@ -0,0 +1,217 @@ +import warnings +from collections.abc import Iterator +from typing import TYPE_CHECKING + +from mmif import Document +from mmif.utils.av_helper import ( + _check_av_dep, + _check_document, + _ffmpeg_op, + _media_check_hint, + _open_container_and_stream, +) +from mmif.utils.timeunit_helper import iter_time_windows +from mmif.vocabulary import DocumentTypes + +if TYPE_CHECKING: + import numpy + +AUDIO_DOCUMENT_TYPES = (DocumentTypes.AudioDocument, DocumentTypes.VideoDocument) +# encoder padding can leave the decoded audio a few tens of ms short of the +# stated duration; anything shorter than this is not reported as truncation +TRUNCATION_TOLERANCE_MS = 100 + + +def iter_audio_windows( + document: Document, + window_ms: int, + stride_ms: int | None = None, + sample_rate: int = 16000, + min_last_window_ms: int = 0, +) -> Iterator[tuple["numpy.ndarray", int, int]]: + """ + Decodes an audio document, or the audio track of a video, into windows. + + Samples are resampled to mono ``float32`` at ``sample_rate`` and decoded + incrementally, so memory use is bounded by one window plus its overlap. + The windows are those of :func:`~mmif.utils.timeunit_helper.iter_time_windows` + from the start of the audio stream, which can be later than zero. After + the last full window, one window runs from the next window start to the + end of the stream, unless it is shorter than ``min_last_window_ms``. + + :param document: audio or video document with an audio stream + :param window_ms: window length in milliseconds + :param stride_ms: distance between window starts; defaults to ``window_ms`` + :param sample_rate: output sample rate in Hz + :param min_last_window_ms: shortest final window to keep + :returns: ``(samples, start_ms, end_ms)`` triples + :raises ValueError: if a size is not positive, ``min_last_window_ms`` is + outside ``[0, window_ms]``, or the document has an unsupported type + :raises FileNotFoundError: if there is no document or no file + :raises RuntimeError: if the file cannot be opened, has no audio stream, + or fails to decode + :warns UserWarning: if the decoded audio ends more than + ``TRUNCATION_TOLERANCE_MS`` before the duration the file states + """ + if window_ms <= 0: + raise ValueError('window_ms must be positive') + if stride_ms is None: + stride_ms = window_ms + elif stride_ms <= 0: + raise ValueError('stride_ms must be positive') + if sample_rate <= 0: + raise ValueError('sample_rate must be positive') + if not 0 <= min_last_window_ms <= window_ms: + raise ValueError('min_last_window_ms must be between 0 and window_ms') + _check_document(document, AUDIO_DOCUMENT_TYPES) + return _iter_audio_windows(document, window_ms, stride_ms, sample_rate, min_last_window_ms) + + +def _iter_audio_windows(document, window_ms, stride_ms, sample_rate, min_last_window_ms): + av = _check_av_dep() + numpy = _check_av_dep('numpy') + container, stream, path = _open_container_and_stream(document, 'audio', AUDIO_DOCUMENT_TYPES) + # channels are averaged here, as librosa does; swresample's own downmix + # scales a stereo pair by the square root of two + resampler = None + stream_start_s = 0.0 + if stream.start_time is not None and stream.time_base is not None: + stream_start_s = float(stream.start_time * stream.time_base) + + # sample indices count from the first decoded sample; `origin_ms` is its + # time on the stream timeline and the start of the window grid + origin_ms = None + origin_s = 0.0 + windows: Iterator[tuple[int, int]] = iter(()) + window = None + chunks = [] # decoded samples from `buffer_start` on, in arrival order + buffer_start = 0 + decoded_end = 0 + window_samples = round(window_ms * sample_rate / 1000) + min_last_samples = max(1, round(min_last_window_ms * sample_rate / 1000)) + + def to_sample(ms): + return round((ms - origin_ms) * sample_rate / 1000) + + def to_ms(sample): + return origin_ms + round(sample * 1000 / sample_rate) + + def append(samples): + nonlocal decoded_end + # a stride longer than the window leaves samples that no window needs + skipped = buffer_start - decoded_end + decoded_end += len(samples) + if skipped > 0: + samples = samples[skipped:] + if len(samples): + chunks.append(samples) + + def buffered(): + if len(chunks) != 1: + chunks[:] = [numpy.concatenate(chunks)] if chunks else [numpy.empty(0, numpy.float32)] + return chunks[0] + + def emit(final=False): + nonlocal window, buffer_start + while True: + if window is None: + window = next(windows) + start_ms, end_ms = window + start, end = to_sample(start_ms), to_sample(end_ms) + if end > decoded_end: + if not final or decoded_end - start < min_last_samples: + return + end, end_ms = decoded_end, to_ms(decoded_end) + yield buffered()[start - buffer_start:end - buffer_start].copy(), start_ms, end_ms + if final: + return + window = None + keep_from = to_sample(start_ms + stride_ms) + chunks[:] = [buffered()[keep_from - buffer_start:]] + buffer_start = keep_from + + def resampled_frames(): + nonlocal resampler + for decoded in container.decode(stream): + if resampler is None: + resampler = av.AudioResampler(format='fltp', layout=decoded.layout, rate=sample_rate) + yield from resampler.resample(decoded) + if resampler is not None: + yield from resampler.resample(None) + + # contiguous resampled frames wait here, so that numpy runs once per + # second of audio and not once per (often 1024-sample) frame + fifo = av.AudioFifo() + + def flush(): + if fifo.samples: + append(fifo.read().to_ndarray().mean(axis=0, dtype=numpy.float32)) + yield from emit() + + stated_end_s = _stated_end_s(av, container, stream, stream_start_s) + + try: + with _ffmpeg_op('decode', document, path, 'audio'): + for frame in resampled_frames(): + frame_s = None + if frame.pts is not None and frame.time_base is not None: + frame_s = float(frame.pts * frame.time_base) + if origin_ms is None: + origin_s = frame_s if frame_s is not None else stream_start_s + origin_ms = round(origin_s * 1000) + windows = iter_time_windows(origin_ms, None, window_ms, stride_ms) + if frame_s is not None: + # a timestamp jump is a gap (filled with silence) or an + # overlap (dropped) on the sample timeline + frame_start = round((frame_s - origin_s) * sample_rate) + gap = frame_start - decoded_end - fifo.samples + if gap: + yield from flush() + while gap > 0: + append(numpy.zeros(min(gap, window_samples), numpy.float32)) + yield from emit() + gap = frame_start - decoded_end + if gap < 0: + samples = frame.to_ndarray().mean(axis=0, dtype=numpy.float32) + append(samples[-gap:]) + yield from emit() + continue + # the fifo checks pts against its own count from zero + frame.pts = None + fifo.write(frame) + if fifo.samples >= sample_rate: + yield from flush() + yield from flush() + if origin_ms is not None: + yield from emit(final=True) + finally: + container.close() + # a cut that ends on a packet boundary decodes cleanly, so only the + # duration the file states shows that audio is missing + decoded_end_ms = to_ms(decoded_end) if origin_ms is not None else round(stream_start_s * 1000) + if stated_end_s is not None: + stated_end_ms = round(stated_end_s * 1000) + if stated_end_ms - decoded_end_ms > TRUNCATION_TOLERANCE_MS: + warnings.warn( + f'The audio of {document.id} ends at {decoded_end_ms} ms, ' + f'{stated_end_ms - decoded_end_ms} ms before the end its file ' + f'states ({stated_end_ms} ms); the windows stop there. ' + f'{_media_check_hint(path, "audio")}', + stacklevel=2) + + +def _stated_end_s(av, container, stream, stream_start_s): + """ + The end of the audio stream according to the file's metadata, in + seconds, or ``None`` when the file does not state it. A container's own + duration counts only when it holds no other stream, since it covers the + longest stream (a video's audio can end before its picture). A WAV file + cut short states its shorter length, since FFmpeg takes it from the file + size, so its truncation cannot be seen here. + """ + if stream.duration is not None and stream.time_base is not None: + return stream_start_s + float(stream.duration * stream.time_base) + if container.duration and len(container.streams) == 1: + start_s = (container.start_time or 0) / av.time_base + return start_s + container.duration / av.time_base + return None diff --git a/mmif/utils/av_helper.py b/mmif/utils/av_helper.py new file mode 100644 index 00000000..5e2f2c94 --- /dev/null +++ b/mmif/utils/av_helper.py @@ -0,0 +1,102 @@ +import importlib +import shlex +from collections.abc import Iterator +from contextlib import contextmanager +from typing import Any + +import mmif +from mmif import Document + + +def _check_av_dep(name: str = 'av') -> Any: + """Imports a package of the ``av`` extra with an actionable message.""" + try: + return importlib.import_module(name) + except ImportError as e: + raise ImportError( + f'Optional package "{e.name}" is not found. ' + 'Install PyAV support with ' + f'`pip install mmif-python[av]=={mmif.__version__}`' + ) from e + + +def _media_check_hint(path: str, media_type: str) -> str: + """Builds an FFmpeg command for checking a media stream.""" + stream_type = {'audio': 'a', 'video': 'v'}.get(media_type, media_type) + return ( + 'The file may be truncated or corrupted; container metadata does not ' + 'reveal this, so to verify the file run ' + f'`ffmpeg -v error -i {shlex.quote(path)} -map 0:{stream_type}:0 ' + '-f null -` and treat any stderr output as a problem ' + '(ffmpeg exits with 0 either way).' + ) + + +@contextmanager +def _ffmpeg_op( + action: str, + document: Document, + path: str, + media_type: str, +) -> Iterator[None]: + """ + Turns failed PyAV operations into actionable ``RuntimeError`` exceptions. + + :param action: operation attempted, as a verb phrase + :param document: MMIF document being read + :param path: resolved local path of the media document + :param media_type: stream type being read, ``"audio"`` or ``"video"`` + :raises RuntimeError: if the operation fails + """ + av = _check_av_dep() + try: + yield + except (av.error.FFmpegError, IndexError) as e: + raise RuntimeError( + f'Cannot {action} {media_type} {document.id} ({path}). ' + f'{_media_check_hint(path, media_type)} ' + f'(original error: {type(e).__name__}: {e})' + ) from e + + +def _check_document(document: Document, document_types: tuple[object, ...]) -> None: + """ + :raises FileNotFoundError: if there is no document + :raises ValueError: if the document has an unsupported type + """ + if document is None: + raise FileNotFoundError('The document does not exist.') + if document.at_type not in document_types: + accepted = ', '.join(str(document_type) for document_type in document_types) + raise ValueError(f'The document must be one of: {accepted}.') + + +def _open_container_and_stream( + document: Document, + media_type: str, + document_types: tuple[object, ...], +) -> tuple[Any, Any, str]: + """ + Opens a media container and selects its first stream of the requested type. + + :param document: MMIF document holding the media file + :param media_type: stream type to select, ``"audio"`` or ``"video"`` + :param document_types: accepted MMIF document types + :returns: open PyAV container, selected stream, and resolved local path + :raises FileNotFoundError: if there is no document or no file + :raises ValueError: if the document has an unsupported type + :raises RuntimeError: if the container or requested stream cannot be read + """ + _check_document(document, document_types) + av = _check_av_dep() + path = str(document.location_path(nonexist_ok=False)) + container = None + with _ffmpeg_op('open', document, path, media_type): + try: + container = av.open(path) + stream = getattr(container.streams, media_type)[0] + except BaseException: + if container is not None: + container.close() + raise + return container, stream, path diff --git a/mmif/utils/timeunit_helper.py b/mmif/utils/timeunit_helper.py index 7f75be37..7732d42a 100644 --- a/mmif/utils/timeunit_helper.py +++ b/mmif/utils/timeunit_helper.py @@ -1,3 +1,4 @@ +from collections.abc import Iterator from datetime import datetime from typing import cast @@ -23,6 +24,54 @@ } +def iter_time_windows( + start_ms: int, + end_ms: int | None, + window_ms: int, + stride_ms: int | None = None, +) -> Iterator[tuple[int, int]]: + """ + Iterates over time windows within a half-open interval. + + The final window is clipped to ``end_ms``. A stride smaller than the + window creates overlapping windows; a larger stride leaves gaps. With + ``end_ms`` of ``None`` the windows do not end, for callers that find the + end of the interval while consuming them. + + :param start_ms: inclusive start of the interval + :param end_ms: exclusive end of the interval, or ``None`` + :param window_ms: window length in milliseconds; must be positive + :param stride_ms: distance between window starts; defaults to + ``window_ms`` and must be positive + :returns: integer-millisecond start/end pairs + :raises ValueError: if the interval is reversed or a size is not positive + """ + if end_ms is not None and end_ms < start_ms: + raise ValueError("end_ms must be greater than or equal to start_ms") + if window_ms <= 0: + raise ValueError("window_ms must be positive") + if stride_ms is None: + stride_ms = window_ms + elif stride_ms <= 0: + raise ValueError("stride_ms must be positive") + return _iter_time_windows(start_ms, end_ms, window_ms, stride_ms) + + +def _iter_time_windows( + start_ms: int, + end_ms: int | None, + window_ms: int, + stride_ms: int, +) -> Iterator[tuple[int, int]]: + window_start_ms = start_ms + while end_ms is None or window_start_ms < end_ms: + window_end_ms = window_start_ms + window_ms + if end_ms is not None: + window_end_ms = min(window_end_ms, end_ms) + yield window_start_ms, window_end_ms + window_start_ms += stride_ms + + def _isoformat_to_millisecond(isoformat: str) -> int: t = datetime.strptime(isoformat, '%H:%M:%S.%f') return int(1000 * (t.hour * 3600 + t.minute * 60 + t.second + t.microsecond / 1000000)) @@ -90,4 +139,3 @@ def convert(t: int | float | str, in_unit: str, out_unit: str, fps: float) -> in # f>i else: return _second_to_isoformat(t / fps) - diff --git a/mmif/utils/video_document_helper.py b/mmif/utils/video_document_helper.py index 9e03fdc8..b560a07f 100644 --- a/mmif/utils/video_document_helper.py +++ b/mmif/utils/video_document_helper.py @@ -1,22 +1,26 @@ import contextvars import importlib import math -import shlex import sys import warnings -from collections.abc import Iterable, Iterator -from collections.abc import Sequence -from typing import cast -from contextlib import contextmanager +from collections.abc import Iterable, Iterator, Sequence from enum import Enum from io import StringIO +from typing import cast import mmif from mmif import Annotation, Document, Mmif, View +from mmif.utils.av_helper import ( + _check_av_dep, + _check_document, + _ffmpeg_op, + _media_check_hint, + _open_container_and_stream, +) from mmif.utils.timeunit_helper import convert from mmif.vocabulary import AnnotationTypes, DocumentTypes -_CV_DEPS = ('av', 'cv2', 'PIL', 'wurlitzer') +_CV_DEPS = ('cv2', 'PIL', 'wurlitzer') _cv_import_warning = ( 'Optional package "{}" is not found. ' 'You might want to install Computer-Vision dependencies ' @@ -29,24 +33,6 @@ 'offset spans (see issue #379).' ) -_MEDIA_CHECK_HINT = ( - 'The file may be truncated or corrupted; container metadata does not ' - 'reveal this, so to verify the file run ' - '`ffmpeg -v error -i {path} -map 0:v:0 -f null -` and treat any stderr ' - 'output as a problem (ffmpeg exits with 0 either way).' -) - - -def _media_check_hint(path: str) -> str: - """ - Formats the "verify this file" hint for a resolved document path. - - :param path: resolved local path of the video document - :returns: hint text, with the path quoted for shell use - :rtype: str - """ - return _MEDIA_CHECK_HINT.format(path=shlex.quote(path)) - def _stream_rate(stream): """ @@ -65,35 +51,6 @@ def _stream_rate(stream): return stream.guessed_rate or stream.average_rate or stream.base_rate -@contextmanager -def _ffmpeg_op(action: str, video_document: Document, path: str): - """ - Turns a failed FFmpeg operation into an actionable ``RuntimeError``. - - Any :py:class:`av.error.FFmpegError` (and the ``IndexError`` of a file - that opens but holds no video stream) is re-raised as a ``RuntimeError`` - naming the document, the media check hint, and the original message. The - original message is repeated in the new one on purpose, because - downstream error reporting (e.g. ``clams-python``'s error view) records - only the outermost exception, so a chained cause would be dropped. - - :param action: what was attempted, as a verb phrase (e.g. ``'open'``) - :param video_document: :py:class:`~mmif.serialize.annotation.Document` - holding the video document being read - :param path: resolved local path of the video document - :raises RuntimeError: if the wrapped operation fails - """ - av = _check_cv_dep('av') - try: - yield - except (av.error.FFmpegError, IndexError) as e: - raise RuntimeError( - f'Cannot {action} video {video_document.id} ({path}). ' - f'{_media_check_hint(path)} ' - f'(original error: {type(e).__name__}: {e})' - ) from e - - def _check_cv_dep(dep): """Import a CV dependency, raising ImportError with a helpful message.""" try: @@ -158,22 +115,9 @@ def open_container(video_document: Document): for instance an mp4 whose ``moov`` atom is cut off, or a prefix too short to hold a video stream """ - av = _check_cv_dep('av') - if video_document is None or video_document.at_type != DocumentTypes.VideoDocument: - raise ValueError(f'The document does not exist.') - - path = cast(str, video_document.location_path(nonexist_ok=False)) - container = None - with _ffmpeg_op('open', video_document, path): - try: - container = av.open(path) - stream = container.streams.video[0] - except BaseException: - # a file can open and still hold no video stream, in which case - # the container must not be left dangling - if container is not None: - container.close() - raise + av = _check_av_dep() + container, stream, _ = _open_container_and_stream( + video_document, 'video', (DocumentTypes.VideoDocument,)) time_base = float(stream.time_base) rate = _stream_rate(stream) fps = round(float(rate), 2) if rate else 0.0 @@ -217,8 +161,7 @@ def get_framerate(video_document: Document) -> float: :param video_document: :py:class:`~mmif.serialize.annotation.Document` instance that holds a video document (``"@type": ".../VideoDocument/..."``) :return: frames per second as a float, rounded to 2 decimal places """ - if video_document is None or video_document.at_type != DocumentTypes.VideoDocument: - raise ValueError(f'The document does not exist.') + _check_document(video_document, (DocumentTypes.VideoDocument,)) framerate_keys = (FPS_DOCPROP_KEY, 'framerate', 'frameRate', 'frame_rate', 'frame-rate', @@ -248,7 +191,7 @@ def _decode_frames(container, stream, decode_errors: list[str]): :param decode_errors: list that collects decoding error texts :returns: generator of :py:class:`av.video.frame.VideoFrame` """ - av = _check_cv_dep('av') + av = _check_av_dep() try: for frame in container.decode(stream): yield frame @@ -299,7 +242,7 @@ def _warn_missing_images( # a zero duration means the container reports none at all, so the # out-of-range test cannot be applied and the file stays suspect if decode_errors or not duration_ms or missing_ms[0] < duration_ms: - msg += f' {_media_check_hint(path)}' + msg += f' {_media_check_hint(path, "video")}' if decode_errors: msg += f' (first decoding error: {decode_errors[0]})' # helper -> public extraction function -> the caller to be pointed at @@ -366,7 +309,7 @@ def _tolerance_ticks(frame): for t_ms in unique_sorted_ms] # seek to the nearest keyframe at or before the earliest target - with _ffmpeg_op('seek in', video_document, path): + with _ffmpeg_op('seek in', video_document, path, media_type='video'): container.seek(target_ticks[0], backward=True, any_frame=False, stream=stream) @@ -875,8 +818,7 @@ def capture(video_document: Document): f'{_PTS_BUG_NOTICE}', DeprecationWarning, stacklevel=2) cv2 = _check_cv_dep('cv2') - if video_document is None or video_document.at_type != DocumentTypes.VideoDocument: - raise ValueError(f'The document does not exist.') + _check_document(video_document, (DocumentTypes.VideoDocument,)) v = cv2.VideoCapture(video_document.location_path(nonexist_ok=False)) fps = round(v.get(cv2.CAP_PROP_FPS), 2) diff --git a/pyproject.toml b/pyproject.toml index 5691624b..8b9a623f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,8 @@ source = "https://github.com/clamsproject/mmif-python" mmif-spec = "https://mmif.clams.ai/1.2.0" [project.optional-dependencies] -cv = ["pillow", "opencv-python", "ffmpeg-python", "wurlitzer", "av>=15"] +av = ["av>=15", "numpy"] +cv = ["mmif-python[av]", "pillow", "opencv-python", "ffmpeg-python", "wurlitzer"] seq = ["numpy"] docs = ["sphinx>=7.0,<8.0", "furo", "m2r2", "autodoc-pydantic"] test = [ diff --git a/tests/test_utils.py b/tests/test_utils.py index 6f7617d5..d5ae0e03 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -17,12 +17,24 @@ DocumentTypes, Mmif ) +from mmif.utils import audio_document_helper as adh +from mmif.utils import av_helper as avh from mmif.utils import sequence_helper as sqh from mmif.utils import text_document_helper as tdh from mmif.utils import timeunit_helper as tuh from mmif.utils import video_document_helper as vdh from mmif.utils import workflow_helper as wfh -from tests import mmif_examples +from tests import mmif_examples + +FIXTURES_DIR = pathlib.Path(__file__).parent + + +def media_document(path, at_type=DocumentTypes.VideoDocument, aid='d1'): + return Document({ + "@type": at_type, + "properties": {"mime": str(at_type).rsplit('/', 2)[-2].lower(), + "id": aid, "location": f"file://{path}"} + }) class TestTimeunitHelper(unittest.TestCase): @@ -46,6 +58,24 @@ def test_convert(self, value): out_t = tuh.convert(converted_back, tu1, 'm', self.FPS) self.assertAlmostEqual(value, out_t, delta=1000/self.FPS) # up to 1 frame error + def test_iter_time_windows(self): + # the last window is clipped to the end of the interval + self.assertEqual([(0, 100), (100, 200), (200, 250)], + list(tuh.iter_time_windows(0, 250, 100))) + # a short stride overlaps, a long one leaves gaps + self.assertEqual([(0, 100), (50, 150), (100, 150)], + list(tuh.iter_time_windows(0, 150, 100, 50))) + self.assertEqual([(0, 100), (300, 400)], + list(tuh.iter_time_windows(0, 400, 100, 300))) + self.assertEqual([], list(tuh.iter_time_windows(5, 5, 100))) + unbounded = tuh.iter_time_windows(10, None, 100) + self.assertEqual([(10, 110), (110, 210)], + [next(unbounded) for _ in range(2)]) + # arguments are checked when called, not on the first `next()` + for args in ((10, 0, 100), (0, 100, 0), (0, 100, 100, 0)): + with self.subTest(args=args), pytest.raises(ValueError): + tuh.iter_time_windows(*args) + @pytest.mark.filterwarnings('ignore::DeprecationWarning') class TestVideoDocumentHelper(unittest.TestCase): @@ -223,12 +253,14 @@ def test_open_container(self): c.close() def test_open_container_rejects_non_video(self): - for bad_doc in (None, Document({ - "@type": "http://mmif.clams.ai/vocabulary/TextDocument/v1", - "properties": {"mime": "text", "id": "t1"}})): - with self.subTest(doc=bad_doc): - with pytest.raises(ValueError): - vdh.open_container(bad_doc) + with pytest.raises(FileNotFoundError): + vdh.open_container(None) + text_doc = media_document('/nonexistent.txt', DocumentTypes.TextDocument) + with pytest.raises(ValueError, match='must be one of'): + vdh.open_container(text_doc) + # get_framerate checks the document the same way + with pytest.raises(ValueError, match='must be one of'): + vdh.get_framerate(text_doc) def test_open_container_missing_file(self): # a missing file is not a broken file, and must stay distinguishable @@ -655,13 +687,7 @@ def setUpClass(cls): cls.tiny = cls._truncate(faststart, tmp / 'tiny.mp4', 200) cls.healthy = healthy_mp4 - @staticmethod - def _doc(path, aid='d1'): - return Document({ - "@type": DocumentTypes.VideoDocument, - "properties": {"mime": "video", "id": aid, - "location": f"file://{path}"} - }) + _doc = staticmethod(media_document) @staticmethod def _remux_faststart(src, dst): @@ -780,7 +806,7 @@ def test_truncated_file_yields_none_past_last_frame(self): # within the tolerance, the last frame itself is reused self.assertEqual(images[1].tobytes(), images[2].tobytes()) self.assertEqual(1, len(messages)) - self.assertIn(vdh._media_check_hint(str(path)), messages[0]) + self.assertIn(avh._media_check_hint(str(path), 'video'), messages[0]) self.assertEqual( expect_decode_error, av.error.InvalidDataError.__name__ in messages[0]) @@ -796,7 +822,7 @@ def test_missing_only_past_duration_is_not_flagged(self): images, messages = self._extract(doc, [duration + 10000]) self.assertEqual([None], images) self.assertEqual(1, len(messages)) - self.assertNotIn(vdh._media_check_hint(str(path)), messages[0]) + self.assertNotIn(avh._media_check_hint(str(path), 'video'), messages[0]) def test_missing_before_and_past_duration_is_flagged(self): # the earliest missing timepoint decides, and both are listed @@ -809,7 +835,7 @@ def test_missing_before_and_past_duration_is_flagged(self): images, messages = self._extract(doc, timepoints) self.assertEqual([None, None], images) self.assertEqual(1, len(messages)) - self.assertIn(vdh._media_check_hint(str(path)), messages[0]) + self.assertIn(avh._media_check_hint(str(path), 'video'), messages[0]) self.assertIn(f'{timepoints[0]}ms, {timepoints[1]}ms', messages[0]) def test_repeated_missing_timepoint_counts_positions(self): @@ -858,7 +884,7 @@ def test_unopenable_files_raise_with_hint(self): with self.subTest(case=label): with pytest.raises(RuntimeError) as excinfo: vdh.open_container(self._doc(path)) - self.assertIn(vdh._media_check_hint(str(path)), + self.assertIn(avh._media_check_hint(str(path), 'video'), str(excinfo.value)) # get_framerate goes through open_container, so it reports the same way with pytest.raises(RuntimeError): @@ -889,7 +915,7 @@ def wrapped(video_document): with mock.patch.object(vdh, 'open_container', wrapped): with pytest.raises(RuntimeError) as excinfo: vdh.extract_images_from_timepoints(doc, [1000]) - self.assertIn(vdh._media_check_hint(str(path)), str(excinfo.value)) + self.assertIn(avh._media_check_hint(str(path), 'video'), str(excinfo.value)) self.assertIn('seek failed', str(excinfo.value)) def test_empty_request_never_opens_the_file(self): @@ -967,6 +993,247 @@ def test_iter_timeframe_images_drops_unreadable(self): self.assertIsNotNone(images[0]) +class TestAudioDocumentHelper(unittest.TestCase): + """ + The fixtures hold a ramp (sample ``i`` has value ``i``), so a window's + first value tells where in the stream it was cut. + """ + + RATE = 8000 + SAMPLES = RATE * 3 + 123 # three full seconds and a 123 ms tail + FULL = 1 / 32768 # int16 to float32 scale + + @classmethod + def setUpClass(cls): + import numpy as np + cls.tmpdir = tempfile.TemporaryDirectory() + cls.addClassCleanup(cls.tmpdir.cleanup) + tmp = pathlib.Path(cls.tmpdir.name) + ramp = np.arange(cls.SAMPLES, dtype='int16')[None, :] + cls.mono = cls._write_pcm(tmp / 'mono.wav', [(ramp, 0)]) + cls.stereo = cls._write_pcm(tmp / 'stereo.wav', + [(np.repeat(ramp, 2, axis=0), 0)]) + # the audio track of a video, starting one second into the timeline + cls.offset = cls._write_pcm(tmp / 'offset.mkv', [(ramp, cls.RATE)], + video=FIXTURES_DIR / 'black-2997fps.mp4') + # one second missing between two packets + cls.gap = cls._write_pcm( + tmp / 'gap.mkv', + [(ramp[:, :cls.RATE], 0), (ramp[:, :cls.RATE], 2 * cls.RATE)]) + # 250 ms missing between two 500 ms packets, shorter than one + # output frame of the resampler + half = cls.RATE // 2 + cls.short_gap = cls._write_pcm( + tmp / 'short_gap.mkv', + [(ramp[:, :half], 0), (ramp[:, :half], half + cls.RATE // 4)]) + cls.video_only = FIXTURES_DIR / 'black-2997fps.mp4' + # copies cut in half: Matroska ends cleanly at the cut, the mp4 + # (index at the front) fails to decode the partial packet + packets = [(np.ascontiguousarray(ramp[:, i:i + cls.RATE // 4]), i) + for i in range(0, cls.SAMPLES, cls.RATE // 4)] + cls.cut_mkv = cls._cut(cls._write_pcm(tmp / 'whole.mkv', packets), + tmp / 'cut.mkv') + cls.cut_mp4 = cls._cut(cls._write_aac_mp4(tmp / 'whole.mp4'), + tmp / 'cut.mp4') + + @classmethod + def _write_pcm(cls, dst, frames, video=None): + """``frames`` are ``(int16 array of shape (channels, n), pts)`` pairs.""" + import av + from fractions import Fraction + with av.open(str(dst), 'w') as out: + if video is not None: + src = av.open(str(video)) + video_in = src.streams.video[0] + video_out = out.add_stream_from_template(video_in) + channels = frames[0][0].shape[0] + layout = 'mono' if channels == 1 else 'stereo' + audio_out = out.add_stream('pcm_s16le', rate=cls.RATE, layout=layout) + for samples, pts in frames: + frame = av.AudioFrame.from_ndarray( + samples, format='s16' if channels == 1 else 's16p', layout=layout) + frame.sample_rate = cls.RATE + frame.pts = pts + frame.time_base = Fraction(1, cls.RATE) + out.mux(audio_out.encode(frame)) + out.mux(audio_out.encode(None)) + if video is not None: + for packet in src.demux(video_in): + if packet.dts is not None: + packet.stream = video_out + out.mux(packet) + src.close() + return dst + + @staticmethod + def _cut(src, dst): + data = pathlib.Path(src).read_bytes() + dst.write_bytes(data[:len(data) // 2]) + return dst + + @classmethod + def _write_aac_mp4(cls, dst, seconds=3): + import av + import numpy as np + from fractions import Fraction + with av.open(str(dst), 'w', options={'movflags': 'faststart'}) as out: + stream = out.add_stream('aac', rate=cls.RATE, layout='mono') + t = np.arange(cls.RATE * seconds) / cls.RATE + tone = (np.sin(2 * np.pi * 440 * t) * 0.3).astype('float32') + frame = av.AudioFrame.from_ndarray(tone[None, :], format='fltp', + layout='mono') + frame.sample_rate = cls.RATE + frame.pts = 0 + frame.time_base = Fraction(1, cls.RATE) + out.mux(stream.encode(frame)) + out.mux(stream.encode(None)) + return dst + + def _windows(self, path, at_type=DocumentTypes.AudioDocument, **kwargs): + kwargs.setdefault('sample_rate', self.RATE) + return list(adh.iter_audio_windows(media_document(path, at_type), **kwargs)) + + def _first_sample(self, samples): + return round(float(samples[0]) / self.FULL) + + def test_windows_cover_the_stream(self): + windows = self._windows(self.mono, window_ms=1000) + self.assertEqual([(0, 1000), (1000, 2000), (2000, 3000), (3000, 3015)], + [(s, e) for _, s, e in windows]) + self.assertTrue(all(isinstance(t, int) for _, s, e in windows for t in (s, e))) + self.assertEqual([self.RATE] * 3 + [123], [len(w) for w, _, _ in windows]) + for samples, start_ms, _ in windows: + self.assertEqual(start_ms * self.RATE // 1000, self._first_sample(samples)) + self.assertEqual('float32', samples.dtype.name) + self.assertEqual(1, samples.ndim) + + def test_stride_and_last_window(self): + # overlapping windows, then one tail from the next window start + windows = self._windows(self.mono, window_ms=1000, stride_ms=500) + self.assertEqual([(0, 1000), (500, 1500), (1000, 2000), (1500, 2500), + (2000, 3000), (2500, 3015)], + [(s, e) for _, s, e in windows]) + self.assertEqual(4123, len(windows[-1][0])) + # a tail shorter than the minimum is dropped + windows = self._windows(self.mono, window_ms=1000, stride_ms=500, + min_last_window_ms=600) + self.assertEqual((2000, 3000), windows[-1][1:]) + # a stride longer than the window skips samples between windows + windows = self._windows(self.mono, window_ms=1000, stride_ms=2000) + self.assertEqual([(0, 1000), (2000, 3000)], [(s, e) for _, s, e in windows]) + self.assertEqual(2 * self.RATE, self._first_sample(windows[1][0])) + + def test_resampling_and_downmix(self): + windows = self._windows(self.stereo, window_ms=1000, sample_rate=16000) + self.assertEqual(4, len(windows)) + self.assertEqual([16000] * 3, [len(w) for w, _, _ in windows[:3]]) + # both channels carry the ramp, so the mono mix is the ramp at twice + # the sample count; the mean of each full window places it + for samples, start_ms, _ in windows[:3]: + expected = (start_ms + 500) * self.RATE / 1000 * self.FULL + self.assertAlmostEqual(expected, float(samples.mean()), delta=self.FULL * 10) + + def test_stream_start_offset(self): + windows = self._windows(self.offset, DocumentTypes.VideoDocument, window_ms=1000) + self.assertEqual([(1000, 2000), (2000, 3000), (3000, 4000), (4000, 4015)], + [(s, e) for _, s, e in windows]) + self.assertEqual(0, self._first_sample(windows[0][0])) + + def test_gap_is_filled_with_silence(self): + windows = self._windows(self.gap, window_ms=1000) + self.assertEqual([(0, 1000), (1000, 2000), (2000, 3000)], + [(s, e) for _, s, e in windows]) + self.assertFalse(windows[1][0].any()) + self.assertEqual([0, 1], [round(float(v) / self.FULL) for v in windows[2][0][:2]]) + + def test_short_gap_while_resampling(self): + # 8 kHz to 16 kHz: the gap falls inside one output frame of the + # resampler, and must still be placed at 500-750 ms + [(samples, start_ms, end_ms)] = self._windows( + self.short_gap, window_ms=2000, sample_rate=16000) + self.assertEqual((0, 1250), (start_ms, end_ms)) + self.assertEqual(20000, len(samples)) + + def at(ms): + return float(samples[ms * 16]) / self.FULL + + # away from the edges, where the resampling filter rings + self.assertAlmostEqual(400 * 8, at(400), delta=2) + self.assertLess(max(abs(at(ms)) for ms in range(550, 700, 10)), 1) + self.assertAlmostEqual(50 * 8, at(800), delta=2) + + def test_rejects_bad_arguments_before_decoding(self): + doc = media_document(self.mono, DocumentTypes.AudioDocument) + for kwargs in (dict(window_ms=0), dict(window_ms=100, stride_ms=0), + dict(window_ms=100, sample_rate=0), + dict(window_ms=100, min_last_window_ms=101)): + with self.subTest(kwargs=kwargs), pytest.raises(ValueError): + adh.iter_audio_windows(doc, **kwargs) + with pytest.raises(FileNotFoundError): + adh.iter_audio_windows(None, window_ms=100) + with pytest.raises(ValueError, match='must be one of'): + adh.iter_audio_windows( + media_document('/nonexistent.txt', DocumentTypes.TextDocument), window_ms=100) + with pytest.raises(FileNotFoundError): + self._windows('/nonexistent.wav', window_ms=100) + + def test_truncated_file_warns_with_hint(self): + # a WAV file cut short states its shorter length, so only a format + # that keeps the full duration in its header can be checked + with pytest.warns(UserWarning, match='before the end') as w: + windows = self._windows(self.cut_mkv, window_ms=1000) + # the windows stop where the decodable audio stops + self.assertLess(windows[-1][2], 2000) + self.assertIn(avh._media_check_hint(str(self.cut_mkv), 'audio'), + str(w[0].message)) + + def test_truncated_mp4_raises_with_hint(self): + with pytest.raises(RuntimeError) as excinfo: + self._windows(self.cut_mp4, window_ms=1000) + self.assertIn(avh._media_check_hint(str(self.cut_mp4), 'audio'), + str(excinfo.value)) + + def test_whole_files_do_not_warn_of_truncation(self): + cases = [(self.mono, DocumentTypes.AudioDocument), + (self.stereo, DocumentTypes.AudioDocument), + (self.gap, DocumentTypes.AudioDocument), + # an audio track inside a longer video + (self.offset, DocumentTypes.VideoDocument)] + for path, at_type in cases: + with self.subTest(fixture=path.name): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + self._windows(path, at_type, window_ms=1000) + self.assertEqual([], [str(w.message) for w in caught + if 'before the end' in str(w.message)]) + + def test_missing_audio_stream_raises_with_hint(self): + with pytest.raises(RuntimeError) as excinfo: + self._windows(self.video_only, DocumentTypes.VideoDocument, window_ms=1000) + self.assertIn(avh._media_check_hint(str(self.video_only), 'audio'), str(excinfo.value)) + self.assertIn('-map 0:a:0', str(excinfo.value)) + + def test_decode_error_raises_with_hint(self): + import av + + class _DecodeFails: + def __init__(self, container): + self._container = container + + def __getattr__(self, name): + return getattr(self._container, name) + + def decode(self, *args, **kwargs): + raise av.error.InvalidDataError(1094995529, 'decode failed') + + real_open = av.open + with mock.patch.object(av, 'open', lambda path: _DecodeFails(real_open(path))): + with pytest.raises(RuntimeError) as excinfo: + self._windows(self.mono, window_ms=1000) + self.assertIn('decode failed', str(excinfo.value)) + self.assertIn(avh._media_check_hint(str(self.mono), 'audio'), str(excinfo.value)) + + class TestSequenceHelper(unittest.TestCase): def test_validate_labelset(self):