adding audio_document_helper in tandem with VDH

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

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

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)
-