diff --git a/av/indexentries.pxd b/av/indexentries.pxd new file mode 100644 index 000000000..44575476a --- /dev/null +++ b/av/indexentries.pxd @@ -0,0 +1,9 @@ +cimport libav as lib + + +cdef class IndexEntries: + cdef lib.AVStream *stream_ptr + cdef _init(self, lib.AVStream *ptr) + + +cdef IndexEntries wrap_index_entries(lib.AVStream *ptr) diff --git a/av/indexentries.py b/av/indexentries.py new file mode 100644 index 000000000..ecbc83d27 --- /dev/null +++ b/av/indexentries.py @@ -0,0 +1,91 @@ +import cython +import cython.cimports.libav as lib +from cython.cimports.av.indexentry import wrap_index_entry +from cython.cimports.libc.stdint import int64_t + +_cinit_bypass_sentinel = cython.declare(object, object()) + + +@cython.cfunc +def wrap_index_entries(ptr: cython.pointer[lib.AVStream]) -> IndexEntries: + obj: IndexEntries = IndexEntries(_cinit_bypass_sentinel) + obj._init(ptr) + return obj + + +@cython.cclass +class IndexEntries: + """A sequence-like view of FFmpeg's per-stream index entries. + + Exposed as :attr:`~av.stream.Stream.index_entries`. + + The index is provided by the demuxer and may be empty or incomplete depending + on the container format. This is useful for fast multi-seek loops (e.g., decoding + at a lower-than-native framerate). + """ + + def __cinit__(self, sentinel): + if sentinel is _cinit_bypass_sentinel: + return + raise RuntimeError("cannot manually instantiate IndexEntries") + + @cython.cfunc + def _init(self, ptr: cython.pointer[lib.AVStream]): + self.stream_ptr = ptr + + def __repr__(self): + return f"" + + def __len__(self) -> int: + with cython.nogil: + return lib.avformat_index_get_entries_count(self.stream_ptr) + + def __iter__(self): + for i in range(len(self)): + yield self[i] + + def __getitem__(self, index): + if isinstance(index, int): + n = len(self) + if index < 0: + index += n + if index < 0 or index >= n: + raise IndexError(f"Index entries {index} out of bounds for size {n}") + + c_idx = cython.declare(cython.int, index) + with cython.nogil: + entry = lib.avformat_index_get_entry(self.stream_ptr, c_idx) + + if entry == cython.NULL: + raise IndexError("index entry not found") + + return wrap_index_entry(entry) + + elif isinstance(index, slice): + start, stop, step = index.indices(len(self)) + return [self[i] for i in range(start, stop, step)] + + else: + raise TypeError("Index must be an integer or a slice") + + def search_timestamp( + self, timestamp, *, backward: bool = True, any_frame: bool = False + ): + """Search the underlying index for ``timestamp``. + + This wraps FFmpeg's ``av_index_search_timestamp``. + + Returns an index into this object, or ``-1`` if no match is found. + """ + c_timestamp = cython.declare(int64_t, timestamp) + flags = cython.declare(cython.int, 0) + + if backward: + flags |= lib.AVSEEK_FLAG_BACKWARD + if any_frame: + flags |= lib.AVSEEK_FLAG_ANY + + with cython.nogil: + idx = lib.av_index_search_timestamp(self.stream_ptr, c_timestamp, flags) + + return idx diff --git a/av/indexentries.pyi b/av/indexentries.pyi new file mode 100644 index 000000000..c5d495711 --- /dev/null +++ b/av/indexentries.pyi @@ -0,0 +1,14 @@ +from typing import Iterator, overload + +from av.indexentry import IndexEntry + +class IndexEntries: + def __len__(self) -> int: ... + def __iter__(self) -> Iterator[IndexEntry]: ... + @overload + def __getitem__(self, index: int) -> IndexEntry: ... + @overload + def __getitem__(self, index: slice) -> list[IndexEntry]: ... + def search_timestamp( + self, timestamp, *, backward: bool = True, any_frame: bool = False + ) -> int: ... diff --git a/av/indexentry.pxd b/av/indexentry.pxd new file mode 100644 index 000000000..55bac3ff4 --- /dev/null +++ b/av/indexentry.pxd @@ -0,0 +1,9 @@ +cimport libav as lib + + +cdef class IndexEntry: + cdef lib.AVIndexEntry *ptr + cdef _init(self, lib.AVIndexEntry *ptr) + + +cdef IndexEntry wrap_index_entry(lib.AVIndexEntry *ptr) diff --git a/av/indexentry.py b/av/indexentry.py new file mode 100644 index 000000000..4d793f0cd --- /dev/null +++ b/av/indexentry.py @@ -0,0 +1,63 @@ +import cython +import cython.cimports.libav as lib + +_cinit_bypass_sentinel = cython.declare(object, object()) + + +@cython.cfunc +def wrap_index_entry(ptr: cython.pointer[lib.AVIndexEntry]) -> IndexEntry: + obj: IndexEntry = IndexEntry(_cinit_bypass_sentinel) + obj._init(ptr) + return obj + + +@cython.cclass +class IndexEntry: + """A single entry from a stream's index. + + This is a thin wrapper around FFmpeg's ``AVIndexEntry``. + + The exact meaning of the fields depends on the container/demuxer. + """ + + def __cinit__(self, sentinel): + if sentinel is not _cinit_bypass_sentinel: + raise RuntimeError("cannot manually instantiate IndexEntry") + + @cython.cfunc + def _init(self, ptr: cython.pointer[lib.AVIndexEntry]): + self.ptr = ptr + + def __repr__(self): + return ( + f"" + ) + + @property + def pos(self): + return self.ptr.pos + + @property + def timestamp(self): + return self.ptr.timestamp + + @property + def flags(self): + return self.ptr.flags + + @property + def is_keyframe(self): + return bool(self.ptr.flags & lib.AVINDEX_KEYFRAME) + + @property + def is_discard(self): + return bool(self.ptr.flags & lib.AVINDEX_DISCARD_FRAME) + + @property + def size(self): + return self.ptr.size + + @property + def min_distance(self): + return self.ptr.min_distance diff --git a/av/indexentry.pyi b/av/indexentry.pyi new file mode 100644 index 000000000..5365bc768 --- /dev/null +++ b/av/indexentry.pyi @@ -0,0 +1,8 @@ +class IndexEntry: + pos: int + timestamp: int + flags: int + is_keyframe: bool + is_discard: bool + size: int + min_distance: int diff --git a/av/stream.pxd b/av/stream.pxd index 9aa6616e5..ecdbe8029 100644 --- a/av/stream.pxd +++ b/av/stream.pxd @@ -3,6 +3,7 @@ cimport libav as lib from av.codec.context cimport CodecContext from av.container.core cimport Container from av.frame cimport Frame +from av.indexentries cimport IndexEntries from av.packet cimport Packet @@ -16,6 +17,8 @@ cdef class Stream: # CodecContext attributes. cdef readonly CodecContext codec_context + cdef readonly IndexEntries index_entries + # Private API. cdef _init(self, Container, lib.AVStream*, CodecContext) cdef _finalize_for_output(self) diff --git a/av/stream.py b/av/stream.py index a43c43496..868def2ab 100644 --- a/av/stream.py +++ b/av/stream.py @@ -3,6 +3,7 @@ import cython from cython.cimports import libav as lib from cython.cimports.av.error import err_check +from cython.cimports.av.indexentries import wrap_index_entries from cython.cimports.av.packet import Packet from cython.cimports.av.utils import ( avdict_to_dict, @@ -106,6 +107,7 @@ def _init( ): self.container = container self.ptr = stream + self.index_entries = wrap_index_entries(self.ptr) self.codec_context = codec_context if self.codec_context: diff --git a/av/stream.pyi b/av/stream.pyi index f6ca196c8..8790d2601 100644 --- a/av/stream.pyi +++ b/av/stream.pyi @@ -4,6 +4,7 @@ from typing import Literal, cast from .codec import Codec, CodecContext from .container import Container +from .indexentries import IndexEntries class Disposition(Flag): default = cast(int, ...) @@ -32,6 +33,7 @@ class Stream: codec: Codec codec_context: CodecContext metadata: dict[str, str] + index_entries: IndexEntries id: int profiles: list[str] profile: str | None diff --git a/include/libavformat/avformat.pxd b/include/libavformat/avformat.pxd index 3a2218f06..1f7ba18f8 100644 --- a/include/libavformat/avformat.pxd +++ b/include/libavformat/avformat.pxd @@ -350,3 +350,18 @@ cdef extern from "libavformat/avformat.h" nogil: # custom cdef set pyav_get_available_formats() + + cdef struct AVIndexEntry: + int64_t pos + int64_t timestamp + int flags + int size + int min_distance + + cdef enum: + AVINDEX_KEYFRAME + AVINDEX_DISCARD_FRAME + + cdef AVIndexEntry *avformat_index_get_entry(AVStream *st, int idx) + cdef int avformat_index_get_entries_count(AVStream *st) + cdef int av_index_search_timestamp(AVStream *st, int64_t timestamp, int flags) diff --git a/tests/test_indexentries.py b/tests/test_indexentries.py new file mode 100644 index 000000000..c0d5cd393 --- /dev/null +++ b/tests/test_indexentries.py @@ -0,0 +1,79 @@ +import av + +from .common import TestCase, fate_suite + + +class TestIndexEntries(TestCase): + def test_index_entries_len_mp4(self) -> None: + with av.open(fate_suite("h264/interlaced_crop.mp4")) as container: + stream = container.streams.video[0] + assert len(stream.index_entries) == stream.frames + + def test_index_entries_len_webm(self) -> None: + with av.open( + fate_suite("vp9-test-vectors/vp90-2-00-quantizer-00.webm") + ) as container: + stream = container.streams.video[0] + index_entries_len_before_demux = len(stream.index_entries) + + keyframes = len([p for p in container.demux(video=0) if p.is_keyframe]) + assert index_entries_len_before_demux == keyframes + + def test_index_entries_search_timestamp_options_mp4(self) -> None: + with av.open(fate_suite("h264/interlaced_crop.mp4")) as container: + stream = container.streams.video[0] + fi = stream.index_entries + + target_ts = -1 + + assert fi.search_timestamp(target_ts) == 0 + assert fi.search_timestamp(target_ts, any_frame=True) == 1 + assert fi.search_timestamp(target_ts, backward=False) == 21 + assert fi.search_timestamp(target_ts, backward=False, any_frame=True) == 1 + + e0 = fi[0] + e1 = fi[1] + e21 = fi[21] + assert e0.timestamp == -2 and e0.is_keyframe + assert e1.timestamp == -1 and not e1.is_keyframe + assert e21.timestamp == 19 and e21.is_keyframe + + def test_index_entries_matches_packet_mp4(self) -> None: + with av.open(fate_suite("h264/interlaced_crop.mp4")) as container: + stream = container.streams.video[0] + for i, packet in enumerate(container.demux(video=0)): + if packet.dts is not None: + assert stream.index_entries[i].timestamp == packet.dts + + def test_index_entries_in_bounds(self) -> None: + with av.open(fate_suite("h264/interlaced_crop.mp4")) as container: + stream = container.streams.video[0] + first = stream.index_entries[0] + first_neg = stream.index_entries[-len(stream.index_entries)] + last = stream.index_entries[-1] + last_pos = stream.index_entries[len(stream.index_entries) - 1] + assert first.timestamp == first_neg.timestamp + assert last.timestamp == last_pos.timestamp + + def test_index_entries_out_of_bounds(self) -> None: + with av.open(fate_suite("h264/interlaced_crop.mp4")) as container: + stream = container.streams.video[0] + with self.assertRaises(IndexError): + _ = stream.index_entries[len(stream.index_entries)] + + with self.assertRaises(IndexError): + _ = stream.index_entries[-len(stream.index_entries) - 1] + + def test_index_entries_slice(self) -> None: + with av.open(fate_suite("h264/interlaced_crop.mp4")) as container: + stream = container.streams.video[0] + + individual_indices = [stream.index_entries[i] for i in range(1, 5)] + slice_indices = stream.index_entries[1:5] + assert len(individual_indices) == len(slice_indices) == 4 + assert all( + [ + i.timestamp == j.timestamp + for i, j in zip(individual_indices, slice_indices) + ] + )