Bitstream filter API fixes

av/bitstream.pxd

@@ -8,3 +8,4 @@ cdef class BitStreamFilterContext:
     cdef const lib.AVBSFContext *ptr
 
     cpdef filter(self, Packet packet=?)
+    cpdef flush(self)

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]

av/bitstream.pyx

@@ -11,27 +11,33 @@ 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
 
         with nogil:
             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()

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
+    )

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)