111 lines
5.0 KiB
Python
111 lines
5.0 KiB
Python
"""One-sample fMP4 → persistent H.264 decoder, no recording/file dependency.
|
|
|
|
Reuses the archive's bounded ISO-BMFF timing validator, not its materializer.
|
|
Only the supplied init and current complete fragment are inspected. Codec state
|
|
persists across fragments; no later fragment or source EOF is supplied to make
|
|
the current frame appear. Other codecs/reordered/multi-sample streams fail
|
|
closed and require a different explicit profile. Run inside a supervised,
|
|
memory-limited CPU child; native decoder allocations are not Python byte counts.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib
|
|
import io
|
|
from fractions import Fraction
|
|
from typing import Any
|
|
|
|
from k1link.media_fragments import (
|
|
ParseBudget,
|
|
single_sample_payload,
|
|
video_fragment_timing,
|
|
video_timing,
|
|
)
|
|
|
|
|
|
class StreamingDecodeError(ValueError):
|
|
pass
|
|
|
|
|
|
class FragmentDecoder:
|
|
def __init__(self) -> None:
|
|
self._av = importlib.import_module("av")
|
|
if self._av.__version__ != "18.0.0":
|
|
raise StreamingDecodeError("unqualified PyAV version")
|
|
# Import before readiness, not on the first source frame's deadline.
|
|
importlib.import_module("numpy")
|
|
self._codec: Any = None
|
|
self._init: bytes | None = None
|
|
self._timing: Any = None
|
|
self._next_dts: int | None = None
|
|
self.frames = 0
|
|
self.failed = False
|
|
|
|
def configure(self, payload: bytes) -> None:
|
|
if self._init is not None or self.failed or not 0 < len(payload) <= 65536:
|
|
self.failed = True
|
|
raise StreamingDecodeError("invalid or repeated codec initialization")
|
|
try:
|
|
timing = video_timing(payload, ParseBudget(128, 1))
|
|
with self._av.open(io.BytesIO(payload), format="mp4", mode="r") as container:
|
|
if len(container.streams) != 1 or len(container.streams.video) != 1:
|
|
raise StreamingDecodeError("one video track is required")
|
|
description = container.streams.video[0].codec_context
|
|
if description.name != "h264" or (description.width, description.height) != (
|
|
800,
|
|
600,
|
|
):
|
|
raise StreamingDecodeError("codec or calibrated native raster changed")
|
|
extra = description.extradata
|
|
if not extra or len(extra) > 65536:
|
|
raise StreamingDecodeError("bounded H264 codec configuration required")
|
|
codec = self._av.CodecContext.create("h264", "r")
|
|
codec.extradata = extra
|
|
codec.thread_count = 1
|
|
codec.thread_type = "SLICE"
|
|
codec.open()
|
|
self._codec, self._init, self._timing = codec, bytes(payload), timing
|
|
except Exception:
|
|
self.failed = True
|
|
raise
|
|
|
|
def decode(self, payload: bytes) -> Any:
|
|
if self.failed or self._init is None or not 0 < len(payload) <= 1024 * 1024:
|
|
self.failed = True
|
|
raise StreamingDecodeError("uninitialized, failed or oversized fragment")
|
|
try:
|
|
timing = video_fragment_timing(payload, self._timing, ParseBudget(128, 1))
|
|
if self.frames == 0 and not timing.random_access:
|
|
raise StreamingDecodeError("epoch must start at a random-access fragment")
|
|
if self._next_dts is not None and timing.base_decode_time != self._next_dts:
|
|
raise StreamingDecodeError("fragment codec timeline is discontinuous")
|
|
# Opening a new demuxer per fragment repeats FFmpeg stream probing.
|
|
# Validate the narrow explicit layout, then feed the one AVCC sample
|
|
# to the persistent codec. No future packet or flush is supplied.
|
|
packet = self._av.Packet(single_sample_payload(payload, self._timing))
|
|
packet.pts = packet.dts = timing.base_decode_time
|
|
packet.duration = timing.duration_units
|
|
packet.time_base = Fraction(1, self._timing.timescale)
|
|
frames = self._codec.decode(packet)
|
|
if len(frames) != 1 or self._codec.has_b_frames:
|
|
raise StreamingDecodeError("current frame requires future input or reordering")
|
|
frame = frames[0]
|
|
if (frame.width, frame.height) != (800, 600) or frame.pts != packet.pts:
|
|
raise StreamingDecodeError("decoded raster or frame identity changed")
|
|
if self.frames == 0 and not frame.key_frame:
|
|
raise StreamingDecodeError("decoder did not confirm the initial keyframe")
|
|
image = frame.to_ndarray(format="bgr24")
|
|
if image.shape != (600, 800, 3) or image.nbytes != 1_440_000:
|
|
raise StreamingDecodeError("invalid native decoder output")
|
|
image.setflags(write=False)
|
|
self._next_dts = timing.base_decode_time + timing.duration_units
|
|
self.frames += 1
|
|
return image
|
|
except Exception:
|
|
self.failed = True
|
|
raise
|
|
|
|
def close(self) -> None:
|
|
self.failed = True
|
|
self._codec = self._init = self._timing = None
|