"""Frame-reading wrapper over torchcodec (infrastructure).
torchcodec and the rest of the PyTorch ecosystem are an **optional** dependency
(the ``media`` extra). Everything here imports them lazily, so importing
:mod:`annie` — and running the bulk of the test suite — never requires a 2 GB
torch install. Call :func:`media_available` to check, and the loaders raise a
clear :class:`MediaUnavailableError` if torchcodec is missing.
Two seek modes are exposed (per the design):
* ``approximate`` — fast preview scrubbing (Browse thumbnails / strips).
* ``exact`` — frame-accurate reads for the Annotator (guarantees frame ``i`` is
frame ``i``).
Frames are returned as ``numpy`` ``uint8`` arrays in HWC/RGB order, which is what
the Pillow-based overlay renderer in :mod:`annie.color` consumes.
"""
from __future__ import annotations
import importlib.util
import sys
from collections.abc import Callable, Generator
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal
from annie.core.runtime import decoder_problem, ensure_windows_ffmpeg_dlls
if TYPE_CHECKING: # pragma: no cover - typing only
import numpy as np
SeekMode = Literal["exact", "approximate"]
"""Frame-seek strategy: ``"exact"`` for frame accuracy, ``"approximate"`` for speed."""
DecoderBackend = Callable[[], Any]
"""Loads the video-decoder class (default: torchcodec's ``VideoDecoder``)."""
_backend: DecoderBackend | None = None
"""The decoder backend in use; ``None`` means torchcodec (see :func:`use_decoder_backend`)."""
[docs]
@contextmanager
def use_decoder_backend(load: DecoderBackend) -> Generator[None]:
"""Decode with another backend while the block runs (e.g. a fake decoder in tests).
The backend replaces torchcodec entirely, including the "is the extra installed"
check; whatever it raises is handled exactly like a torchcodec load failure.
Args:
load: Returns the decoder class, or raises ``ImportError``/``OSError`` when the
backend cannot load.
"""
global _backend
previous, _backend = _backend, load
try:
yield
finally:
_backend = previous
def _require_media() -> None:
"""Raise :class:`MediaUnavailableError` if torchcodec is not installed."""
if not media_available():
raise MediaUnavailableError(
"Frame decoding requires the 'media' extra. Install it with "
'`uv pip install -e ".[media]"` and a system FFmpeg (4-8).'
)
#: Why the installed decoder cannot load, once :func:`check_decoder` has found out.
_decoder_error: MediaUnavailableError | None = None
def _video_decoder_class(): # noqa: ANN202 - external type
"""Import torchcodec's ``VideoDecoder``, turning a load failure into a clear error.
torchcodec can be installed and still fail to import, e.g. when PyTorch's DLLs miss
the Visual C++ runtime on Windows. That surfaces as an ``OSError`` deep inside
``import torch``; callers get a :class:`MediaUnavailableError` saying what to do.
"""
if _backend is None:
_require_media()
ensure_windows_ffmpeg_dlls() # Windows: FFmpeg's DLLs must be registered before import
try:
if _backend is not None:
return _backend()
from torchcodec.decoders import VideoDecoder # local import: optional dependency
except (ImportError, OSError) as exc:
raise MediaUnavailableError(decoder_problem(exc)) from exc
return VideoDecoder
[docs]
def check_decoder() -> MediaUnavailableError | None:
"""Load the decoder once and remember why it failed (run off the event loop).
Returns:
The error, whose message says what to do, or ``None`` when frames can be decoded
or the ``media`` extra is not installed (which is reported separately).
"""
global _decoder_error
if _backend is None and not media_available():
return None
try:
# torchcodec loads FFmpeg (its core library) only when a decoder is created, and
# does so before opening the file: a path that does not exist shows whether that
# load works, without needing a video.
_decoder(Path(__file__).with_name("no-such-video.mp4"), "approximate")
except MediaUnavailableError as exc:
_decoder_error = exc
except Exception: # noqa: BLE001, S110 - "file not found" etc.: the decoder itself works
pass
return _decoder_error
[docs]
def decoder_problem_found() -> str | None:
"""The problem :func:`check_decoder` found, without loading anything."""
return None if _decoder_error is None else str(_decoder_error)
def _decoder(path: str | Path, seek_mode: SeekMode): # noqa: ANN202 - external type
"""Construct a torchcodec ``VideoDecoder`` for ``path`` in the given seek mode.
Raises:
MediaUnavailableError: If torchcodec cannot load FFmpeg (``Could not load
libtorchcodec``): a problem of the decoder, not of this video, so callers
report it once instead of per video.
"""
decoder_class = _video_decoder_class()
try:
return decoder_class(str(Path(path)), seek_mode=seek_mode)
except RuntimeError as exc:
if "libtorchcodec" in str(exc):
raise MediaUnavailableError(decoder_problem(exc)) from exc
raise
def _to_hwc_uint8(frame) -> np.ndarray: # noqa: ANN001 - external tensor type
"""Convert a torchcodec CHW uint8 tensor to an HWC uint8 numpy array."""
return frame.permute(1, 2, 0).contiguous().cpu().numpy()
[docs]
def read_frame(path: str | Path, frame_idx: int, *, seek_mode: SeekMode = "exact") -> np.ndarray:
"""Read a single frame by index.
Args:
path: Path to the video file.
frame_idx: Zero-based frame index to read.
seek_mode: ``"exact"`` (frame-accurate) or ``"approximate"`` (fast).
Returns:
The frame as an HWC ``uint8`` RGB numpy array.
Raises:
MediaUnavailableError: If torchcodec is not installed.
"""
decoder = _decoder(path, seek_mode)
return _to_hwc_uint8(decoder[frame_idx])
[docs]
def read_frames(
path: str | Path, indices: list[int], *, seek_mode: SeekMode = "approximate"
) -> list[np.ndarray]:
"""Read several frames by index.
Args:
path: Path to the video file.
indices: Zero-based frame indices to read.
seek_mode: Seek mode; defaults to ``"approximate"`` for fast previews.
Returns:
One HWC ``uint8`` RGB numpy array per requested index, in order.
Raises:
MediaUnavailableError: If torchcodec is not installed.
"""
decoder = _decoder(path, seek_mode)
return [_to_hwc_uint8(decoder[i]) for i in indices]
[docs]
def frame_count(path: str | Path) -> int:
"""Return the number of frames in a video.
Args:
path: Path to the video file.
Returns:
The total frame count.
Raises:
MediaUnavailableError: If torchcodec is not installed.
"""
decoder = _decoder(path, "approximate")
return int(decoder.metadata.num_frames)
[docs]
def read_strip(
path: str | Path, count: int = 5, *, seek_mode: SeekMode = "exact"
) -> tuple[list[int], list[np.ndarray], int]:
"""Read evenly spaced sample frames for the Browse strip in one pass.
Defaults to ``"exact"`` seek mode: in ``"approximate"`` mode torchcodec derives
``num_frames`` from duration × fps, which can over-report the true count for
some clips and make the last sample index point past the real end (raising
"no more frames left to decode"). Exact mode scans the file so the count and
every sampled index are valid; decoding is still guarded per frame.
Args:
path: Path to the video file.
count: Number of evenly spaced frames (first / ¼ / ½ / ¾ / last by default).
seek_mode: Seek mode; ``"exact"`` for an accurate frame count.
Returns:
An ``(indices, frames, num_frames)`` triple: the sampled frame indices that
decoded, their HWC ``uint8`` RGB arrays (aligned), and the total frame count.
Raises:
MediaUnavailableError: If torchcodec is not installed.
"""
decoder = _decoder(path, seek_mode)
total = int(decoder.metadata.num_frames)
indices: list[int] = []
frames: list[np.ndarray] = []
for raw in strip_indices(total, count):
idx = min(raw, total - 1)
try:
frames.append(_to_hwc_uint8(decoder[idx]))
indices.append(idx)
except (RuntimeError, IndexError):
break # tolerate a frame-count overshoot at the tail
return indices, frames, total
[docs]
def strip_indices(total: int, count: int = 5) -> list[int]:
"""Compute evenly spaced sample indices for the Browse 5-frame strip.
Returns indices at first / ¼ / ½ / ¾ / last by default. Pure arithmetic, so
it needs no media backend and is unit-tested directly.
Args:
total: Total number of frames in the video (must be ``>= 1``).
count: How many sample indices to return (``>= 1``).
Returns:
A sorted list of ``count`` distinct-where-possible frame indices.
Raises:
ValueError: If ``total < 1`` or ``count < 1``.
"""
if total < 1:
raise ValueError("total must be >= 1")
if count < 1:
raise ValueError("count must be >= 1")
if count == 1:
return [0]
last = total - 1
return [round(i * last / (count - 1)) for i in range(count)]