From 82e3397e0a42b559118821368c8f31af95829db8 Mon Sep 17 00:00:00 2001 From: Santtu Keskinen Date: Wed, 10 Apr 2024 01:44:49 +0300 Subject: [PATCH] Add bitstream filters --------- Co-authored-by: WyattBlue --- av/__init__.py | 1 + av/bitstream.pxd | 10 +++++ av/bitstream.pyi | 10 +++++ av/bitstream.pyx | 82 ++++++++++++++++++++++++++++++++++++++ av/packet.pyi | 6 +-- docs/api/bitstream.rst | 9 +++++ docs/index.rst | 1 + include/libav.pxd | 1 + include/libavcodec/bsf.pxd | 36 +++++++++++++++++ tests/test_bitstream.py | 62 ++++++++++++++++++++++++++++ 10 files changed, 215 insertions(+), 3 deletions(-) create mode 100644 av/bitstream.pxd create mode 100644 av/bitstream.pyi create mode 100644 av/bitstream.pyx create mode 100644 docs/api/bitstream.rst create mode 100644 include/libavcodec/bsf.pxd create mode 100644 tests/test_bitstream.py diff --git a/av/__init__.py b/av/__init__.py index 8c87e17cc..b4705dc69 100644 --- a/av/__init__.py +++ b/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 diff --git a/av/bitstream.pxd b/av/bitstream.pxd new file mode 100644 index 000000000..2ebc1da75 --- /dev/null +++ b/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=?) diff --git a/av/bitstream.pyi b/av/bitstream.pyi new file mode 100644 index 000000000..c3d4232cd --- /dev/null +++ b/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] diff --git a/av/bitstream.pyx b/av/bitstream.pyx new file mode 100644 index 000000000..b1fedb761 --- /dev/null +++ b/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() diff --git a/av/packet.pyi b/av/packet.pyi index cca33009c..91775eff7 100644 --- a/av/packet.pyi +++ b/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]: ... diff --git a/docs/api/bitstream.rst b/docs/api/bitstream.rst new file mode 100644 index 000000000..4dab1bde5 --- /dev/null +++ b/docs/api/bitstream.rst @@ -0,0 +1,9 @@ + +Bitstream Filters +================= + +.. automodule:: av.bitstream + + .. autoclass:: BitStreamFilterContext + :members: + diff --git a/docs/index.rst b/docs/index.rst index a1a6d3dff..66c8ffbba 100644 --- a/docs/index.rst +++ b/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>`; diff --git a/include/libav.pxd b/include/libav.pxd index b9bfe3943..4312be2e8 100644 --- a/include/libav.pxd +++ b/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" diff --git a/include/libavcodec/bsf.pxd b/include/libavcodec/bsf.pxd new file mode 100644 index 000000000..27619442a --- /dev/null +++ b/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 + ) diff --git a/tests/test_bitstream.py b/tests/test_bitstream.py new file mode 100644 index 000000000..ddc1bfde2 --- /dev/null +++ b/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))