From 21ddc604fd71d6c2c85c23af8f9ef8780c493f57 Mon Sep 17 00:00:00 2001 From: Santtu Keskinen Date: Wed, 17 Apr 2024 19:53:56 +0300 Subject: [PATCH] Bitstream filter API fixes --- av/bitstream.pxd | 1 + av/bitstream.pyi | 6 +++++- av/bitstream.pyx | 25 ++++++++++++++++------ include/libavcodec/bsf.pxd | 4 ++++ tests/test_bitstream.py | 44 ++++++++++++++++++++++++++++++++++---- 5 files changed, 69 insertions(+), 11 deletions(-) diff --git a/av/bitstream.pxd b/av/bitstream.pxd index 2ebc1da75..620fa0366 100644 --- a/av/bitstream.pxd +++ b/av/bitstream.pxd @@ -8,3 +8,4 @@ cdef class BitStreamFilterContext: cdef const lib.AVBSFContext *ptr cpdef filter(self, Packet packet=?) + cpdef flush(self) diff --git a/av/bitstream.pyi b/av/bitstream.pyi index c3d4232cd..477c65f2d 100644 --- a/av/bitstream.pyi +++ b/av/bitstream.pyi @@ -3,8 +3,12 @@ from .stream import Stream class BitStreamFilterContext: def __init__( - self, filter_description: str | bytes, stream: Stream | None = None + self, + filter_description: str | bytes, + in_stream: Stream | None = None, + out_stream: Stream | None = None, ): ... def filter(self, packet: Packet | None) -> list[Packet]: ... + def flush(self) -> None: ... bitstream_filters_available: set[str] diff --git a/av/bitstream.pyx b/av/bitstream.pyx index b1fedb761..b5361f8c2 100644 --- a/av/bitstream.pyx +++ b/av/bitstream.pyx @@ -11,8 +11,11 @@ cdef class BitStreamFilterContext: Initializes a bitstream filter: a way to directly modify packet data. Wraps :ffmpeg:`AVBSFContext` + + :param Stream in_stream: A stream that defines the input codec for the bitfilter. + :param Stream out_stream: A stream whose codec is overwritten using the output parameters from the bitfilter. """ - def __cinit__(self, filter_description, Stream stream=None): + def __cinit__(self, filter_description, Stream in_stream=None, Stream out_stream=None): cdef int res cdef char *filter_str = filter_description @@ -20,18 +23,21 @@ cdef class BitStreamFilterContext: res = lib.av_bsf_list_parse_str(filter_str, &self.ptr) err_check(res) - if stream is not None: - with nogil: - res = lib.avcodec_parameters_copy(self.ptr.par_in, stream.ptr.codecpar) - err_check(res) + if in_stream is not None: with nogil: - res = lib.avcodec_parameters_copy(self.ptr.par_out, stream.ptr.codecpar) + res = lib.avcodec_parameters_copy(self.ptr.par_in, in_stream.ptr.codecpar) err_check(res) with nogil: res = lib.av_bsf_init(self.ptr) err_check(res) + if out_stream is not None: + with nogil: + res = lib.avcodec_parameters_copy(out_stream.ptr.codecpar, self.ptr.par_out) + err_check(res) + lib.avcodec_parameters_to_context(out_stream.codec_context.ptr, out_stream.ptr.codecpar) + def __dealloc__(self): if self.ptr: lib.av_bsf_free(&self.ptr) @@ -65,6 +71,13 @@ cdef class BitStreamFilterContext: output.append(new_packet) + cpdef flush(self): + """ + Reset the internal state of the filter. + Should be called e.g. when seeking. + Can be used to make the filter usable again after draining it with EOF marker packet. + """ + lib.av_bsf_flush(self.ptr) cdef get_filter_names(): names = set() diff --git a/include/libavcodec/bsf.pxd b/include/libavcodec/bsf.pxd index 27619442a..4a558b478 100644 --- a/include/libavcodec/bsf.pxd +++ b/include/libavcodec/bsf.pxd @@ -34,3 +34,7 @@ cdef extern from "libavcodec/bsf.h" nogil: AVBSFContext *ctx, AVPacket *pkt ) + + cdef void av_bsf_flush( + AVBSFContext *ctx + ) diff --git a/tests/test_bitstream.py b/tests/test_bitstream.py index ddc1bfde2..937448fe5 100644 --- a/tests/test_bitstream.py +++ b/tests/test_bitstream.py @@ -5,6 +5,11 @@ from .common import TestCase, fate_suite +def is_annexb(packet: Packet) -> bool: + data = bytes(packet) + return data[:3] == b"\0\0\x01" or data[:4] == b"\0\0\0\x01" + + class TestBitStreamFilters(TestCase): def test_filters_availible(self) -> None: self.assertIn("h264_mp4toannexb", bitstream_filters_available) @@ -43,10 +48,6 @@ def test_filter_setts(self) -> None: self.assertEqual(result_packets[1].pts, 1) def test_filter_h264_mp4toannexb(self) -> None: - def is_annexb(packet: Packet) -> bool: - data = bytes(packet) - return data[:3] == b"\0\0\x01" or data[:4] == b"\0\0\0\x01" - with av.open(fate_suite("h264/interlaced_crop.mp4"), "r") as container: stream = container.streams.video[0] ctx = BitStreamFilterContext("h264_mp4toannexb", stream) @@ -60,3 +61,38 @@ def is_annexb(packet: Packet) -> bool: for p in res_packets: self.assertTrue(is_annexb(p)) + + def test_filter_output_parameters(self) -> None: + with av.open(fate_suite("h264/interlaced_crop.mp4"), "r") as container: + stream = container.streams.video[0] + + self.assertFalse(is_annexb(stream.codec_context.extradata)) + ctx = BitStreamFilterContext("h264_mp4toannexb", stream) + self.assertFalse(is_annexb(stream.codec_context.extradata)) + del ctx + + _ = BitStreamFilterContext("h264_mp4toannexb", stream, out_stream=stream) + self.assertTrue(is_annexb(stream.codec_context.extradata)) + + def test_filter_flush(self) -> None: + with av.open(fate_suite("h264/interlaced_crop.mp4"), "r") as container: + stream = container.streams.video[0] + ctx = BitStreamFilterContext("h264_mp4toannexb", stream) + + res_packets = [] + for p in container.demux(stream): + res_packets.extend(ctx.filter(p)) + self.assertEqual(len(res_packets), stream.frames) + + container.seek(0) + # Without flushing, we expect to get an error: "A non-NULL packet sent after an EOF." + with self.assertRaises(ValueError): + for p in container.demux(stream): + ctx.filter(p) + + ctx.flush() + container.seek(0) + for p in container.demux(stream): + res_packets.extend(ctx.filter(p)) + + self.assertEqual(len(res_packets), stream.frames * 2)