Add bitstream filters

---------

Co-authored-by: WyattBlue <[email protected]>

av/__init__.py

@@ -29,6 +29,7 @@
 from av.audio.frame import AudioFrame
 from av.audio.layout import AudioLayout
 from av.audio.resampler import AudioResampler
+from av.bitstream import BitStreamFilterContext, bitstream_filters_available
 from av.codec.codec import Codec, codecs_available
 from av.codec.context import CodecContext
 from av.container import open

av/bitstream.pxd

@@ -0,0 +1,10 @@
+cimport libav as lib
+
+from av.packet cimport Packet
+
+
+cdef class BitStreamFilterContext:
+
+    cdef const lib.AVBSFContext *ptr
+
+    cpdef filter(self, Packet packet=?)

av/bitstream.pyi

@@ -0,0 +1,10 @@
+from .packet import Packet
+from .stream import Stream
+
+class BitStreamFilterContext:
+    def __init__(
+        self, filter_description: str | bytes, stream: Stream | None = None
+    ): ...
+    def filter(self, packet: Packet | None) -> list[Packet]: ...
+
+bitstream_filters_available: set[str]

av/bitstream.pyx

@@ -0,0 +1,82 @@
+cimport libav as lib
+from libc.errno cimport EAGAIN
+
+from av.error cimport err_check
+from av.packet cimport Packet
+from av.stream cimport Stream
+
+
+cdef class BitStreamFilterContext:
+    """
+    Initializes a bitstream filter: a way to directly modify packet data.
+
+    Wraps :ffmpeg:`AVBSFContext`
+    """
+    def __cinit__(self, filter_description, Stream 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)
+            with nogil:
+                res = lib.avcodec_parameters_copy(self.ptr.par_out, stream.ptr.codecpar)
+            err_check(res)
+
+        with nogil:
+            res = lib.av_bsf_init(self.ptr)
+        err_check(res)
+
+    def __dealloc__(self):
+        if self.ptr:
+            lib.av_bsf_free(&self.ptr)
+
+    cpdef filter(self, Packet packet=None):
+        """
+        Processes a packet based on the filter_description set during initialization.
+        Multiple packets may be created.
+
+        :type: list[Packet]
+        """
+        cdef int res
+        cdef Packet new_packet
+
+        with nogil:
+            res = lib.av_bsf_send_packet(self.ptr, packet.ptr if packet is not None else NULL)
+        err_check(res)
+
+        output = []
+        while True:
+            new_packet = Packet()
+            with nogil:
+                res = lib.av_bsf_receive_packet(self.ptr, new_packet.ptr)
+
+            if res == -EAGAIN or res == lib.AVERROR_EOF:
+                return output
+
+            err_check(res)
+            if res:
+                return output
+
+            output.append(new_packet)
+
+
+cdef get_filter_names():
+    names = set()
+    cdef const lib.AVBitStreamFilter *ptr
+    cdef void *opaque = NULL
+    while True:
+        ptr = lib.av_bsf_iterate(&opaque)
+        if ptr:
+            names.add(ptr.name)
+        else:
+            break
+
+    return names
+
+bitstream_filters_available = get_filter_names()

av/packet.pyi

@@ -1,11 +1,11 @@
 from fractions import Fraction
-from typing import Iterator
+from typing import Buffer, Iterator
 
 from av.subtitles.subtitle import SubtitleSet
 
 from .stream import Stream
 
-class Packet:
+class Packet(Buffer):
     stream: Stream
     stream_index: int
     time_base: Fraction
@@ -20,5 +20,5 @@ class Packet:
     is_trusted: bool
     is_disposable: bool
 
-    def __init__(self, input: int | None = None) -> None: ...
+    def __init__(self, input: int | bytes | None = None) -> None: ...
     def decode(self) -> Iterator[SubtitleSet]: ...

docs/api/bitstream.rst

@@ -0,0 +1,9 @@
+
+Bitstream Filters
+=================
+
+.. automodule:: av.bitstream
+
+    .. autoclass:: BitStreamFilterContext
+        :members:
+

docs/index.rst

@@ -21,6 +21,7 @@ Currently we provide:
 - ``libavcodec``:
   :class:`.Codec`,
   :class:`.CodecContext`,
+  :class:`.BitStreamFilterContext`,
   audio/video :class:`frames <.Frame>`,
   :class:`data planes <.Plane>`,
   :class:`subtitles <.Subtitle>`;

include/libav.pxd

@@ -7,6 +7,7 @@ include "libavutil/samplefmt.pxd"
 include "libavutil/motion_vector.pxd"
 
 include "libavcodec/avcodec.pxd"
+include "libavcodec/bsf.pxd"
 include "libavdevice/avdevice.pxd"
 include "libavformat/avformat.pxd"
 include "libswresample/swresample.pxd"

include/libavcodec/bsf.pxd

@@ -0,0 +1,36 @@
+
+cdef extern from "libavcodec/bsf.h" nogil:
+
+    cdef struct AVBitStreamFilter:
+        const char *name
+        AVCodecID *codec_ids
+
+    cdef struct AVCodecParameters:
+        pass
+
+    cdef struct AVBSFContext:
+        const AVBitStreamFilter *filter
+        const AVCodecParameters *par_in
+        const AVCodecParameters *par_out
+
+    cdef const AVBitStreamFilter* av_bsf_get_by_name(const char *name)
+
+    cdef int av_bsf_list_parse_str(
+        const char *str,
+        AVBSFContext **bsf
+    )
+
+    cdef int av_bsf_init(AVBSFContext *ctx)
+    cdef void av_bsf_free(AVBSFContext **ctx)
+
+    cdef AVBitStreamFilter* av_bsf_iterate(void **opaque)
+
+    cdef int av_bsf_send_packet(
+        AVBSFContext *ctx,
+        AVPacket *pkt
+    )
+
+    cdef int av_bsf_receive_packet(
+        AVBSFContext *ctx,
+        AVPacket *pkt
+    )

tests/test_bitstream.py

@@ -0,0 +1,62 @@
+import av
+from av import Packet
+from av.bitstream import BitStreamFilterContext, bitstream_filters_available
+
+from .common import TestCase, fate_suite
+
+
+class TestBitStreamFilters(TestCase):
+    def test_filters_availible(self) -> None:
+        self.assertIn("h264_mp4toannexb", bitstream_filters_available)
+
+    def test_filter_chomp(self) -> None:
+        ctx = BitStreamFilterContext("chomp")
+
+        src_packets: tuple[Packet, None] = (Packet(b"\x0012345\0\0\0"), None)
+        self.assertEqual(bytes(src_packets[0]), b"\x0012345\0\0\0")
+
+        result_packets = []
+        for p in src_packets:
+            result_packets.extend(ctx.filter(p))
+
+        self.assertEqual(len(result_packets), 1)
+        self.assertEqual(bytes(result_packets[0]), b"\x0012345")
+
+    def test_filter_setts(self) -> None:
+        ctx = BitStreamFilterContext("setts=pts=N")
+
+        ctx2 = BitStreamFilterContext(b"setts=pts=N")
+        del ctx2
+
+        p1 = Packet(b"\0")
+        p1.pts = 42
+        p2 = Packet(b"\0")
+        p2.pts = 50
+        src_packets = [p1, p2, None]
+
+        result_packets: list[Packet] = []
+        for p in src_packets:
+            result_packets.extend(ctx.filter(p))
+
+        self.assertEqual(len(result_packets), 2)
+        self.assertEqual(result_packets[0].pts, 0)
+        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)
+
+            res_packets = []
+            for p in container.demux(stream):
+                self.assertFalse(is_annexb(p))
+                res_packets.extend(ctx.filter(p))
+
+            self.assertEqual(len(res_packets), stream.frames)
+
+            for p in res_packets:
+                self.assertTrue(is_annexb(p))