48 lines
1.7 KiB
Python
48 lines
1.7 KiB
Python
"""Decoder selection independent of concrete codec libraries."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Iterable, Iterator
|
|
from itertools import chain
|
|
|
|
from yovision_brain.input import CancellationToken, InputPacket
|
|
|
|
from .models import DecodedFrame, DecoderBackend, DecoderError
|
|
from .raw_rgb import RawRGBDecoder
|
|
from .y4m import Y4MDecoder
|
|
|
|
|
|
class DecoderPipeline:
|
|
def __init__(self, backends: Iterable[DecoderBackend] | None = None) -> None:
|
|
selected = tuple(backends) if backends is not None else (RawRGBDecoder(), Y4MDecoder())
|
|
self._backends: dict[str, DecoderBackend] = {}
|
|
for backend in selected:
|
|
for media_format in backend.media_formats:
|
|
if media_format in self._backends:
|
|
raise ValueError(f"duplicate decoder for media format {media_format!r}")
|
|
self._backends[media_format] = backend
|
|
|
|
def decode(
|
|
self,
|
|
packets: Iterable[InputPacket],
|
|
cancellation: CancellationToken | None = None,
|
|
) -> Iterator[DecodedFrame]:
|
|
iterator = iter(packets)
|
|
if cancellation is not None and cancellation.cancelled:
|
|
return
|
|
try:
|
|
first = next(iterator)
|
|
except StopIteration:
|
|
return
|
|
backend = self._backends.get(first.media_format)
|
|
if backend is None:
|
|
raise DecoderError(f"no decoder registered for media format {first.media_format!r}")
|
|
yield from backend.decode(chain((first,), iterator), cancellation)
|
|
|
|
|
|
def decode_packets(
|
|
packets: Iterable[InputPacket],
|
|
cancellation: CancellationToken | None = None,
|
|
) -> Iterator[DecodedFrame]:
|
|
return DecoderPipeline().decode(packets, cancellation)
|