From 9b2864f301f0c658d793102458cea1c252c64b91 Mon Sep 17 00:00:00 2001 From: Santtu Keskinen Date: Mon, 8 Apr 2024 22:33:41 +0300 Subject: [PATCH 1/4] Add bitstream filters --- av/__init__.py | 1 + av/bitstream.pxd | 9 +++++ av/bitstream.pyx | 77 ++++++++++++++++++++++++++++++++++++++ include/libav.pxd | 1 + include/libavcodec/bsf.pxd | 36 ++++++++++++++++++ tests/test_bitstream.py | 60 +++++++++++++++++++++++++++++ 6 files changed, 184 insertions(+) create mode 100644 av/bitstream.pxd create mode 100644 av/bitstream.pyx 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..ae88568c1 --- /dev/null +++ b/av/bitstream.pxd @@ -0,0 +1,9 @@ +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.pyx b/av/bitstream.pyx new file mode 100644 index 000000000..af04f2e01 --- /dev/null +++ b/av/bitstream.pyx @@ -0,0 +1,77 @@ +from libc.errno cimport EAGAIN + +cimport libav as lib + +from av.error cimport err_check +from av.packet cimport Packet +from av.stream cimport Stream + + +cdef class BitStreamFilterContext: + + 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) + + def _send(self, Packet packet=None): + cdef int res + with nogil: + res = lib.av_bsf_send_packet(self.ptr, packet.ptr if packet is not None else NULL) + err_check(res) + + def _recv(self): + cdef Packet packet = Packet() + + cdef int res + with nogil: + res = lib.av_bsf_receive_packet(self.ptr, packet.ptr) + if res == -EAGAIN or res == lib.AVERROR_EOF: + return + err_check(res) + + if not res: + return packet + + cpdef filter(self, Packet packet=None): + self._send(packet) + + output = [] + while True: + packet = self._recv() + if packet: + output.append(packet) + else: + return output + + +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/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..6fd930772 --- /dev/null +++ b/tests/test_bitstream.py @@ -0,0 +1,60 @@ +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): + self.assertIn('h264_mp4toannexb', bitstream_filters_available) + + def test_filter_chomp(self): + ctx = BitStreamFilterContext('chomp') + + src_packets = [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): + ctx = BitStreamFilterContext('setts=pts=N') + + p1 = Packet(b'\0') + p1.pts = 42 + p2 = Packet(b'\0') + p2.pts = 50 + src_packets = [p1, p2, None] + + result_packets = [] + 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): + def is_annexb(packet): + 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)) From 9a161321cd128ea5b4a91a6df14b78693ba91c3e Mon Sep 17 00:00:00 2001 From: Santtu Keskinen Date: Tue, 9 Apr 2024 10:45:54 +0300 Subject: [PATCH 2/4] Change ' -> " to appease linter --- tests/test_bitstream.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/tests/test_bitstream.py b/tests/test_bitstream.py index 6fd930772..da184ffa8 100644 --- a/tests/test_bitstream.py +++ b/tests/test_bitstream.py @@ -8,27 +8,27 @@ class TestBitStreamFilters(TestCase): def test_filters_availible(self): - self.assertIn('h264_mp4toannexb', bitstream_filters_available) + self.assertIn("h264_mp4toannexb", bitstream_filters_available) def test_filter_chomp(self): - ctx = BitStreamFilterContext('chomp') + ctx = BitStreamFilterContext("chomp") - src_packets = [Packet(b'\x0012345\0\0\0'), None] - self.assertEqual(bytes(src_packets[0]), b'\x0012345\0\0\0') + src_packets = [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') + self.assertEqual(bytes(result_packets[0]), b"\x0012345") def test_filter_setts(self): - ctx = BitStreamFilterContext('setts=pts=N') + ctx = BitStreamFilterContext("setts=pts=N") - p1 = Packet(b'\0') + p1 = Packet(b"\0") p1.pts = 42 - p2 = Packet(b'\0') + p2 = Packet(b"\0") p2.pts = 50 src_packets = [p1, p2, None] @@ -43,11 +43,11 @@ def test_filter_setts(self): def test_filter_h264_mp4toannexb(self): def is_annexb(packet): data = bytes(packet) - return data[:3] == b'\0\0\x01' or data[:4] == b'\0\0\0\x01' + 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: + with av.open(fate_suite("h264/interlaced_crop.mp4"), "r") as container: stream = container.streams.video[0] - ctx = BitStreamFilterContext('h264_mp4toannexb', stream) + ctx = BitStreamFilterContext("h264_mp4toannexb", stream) res_packets = [] for p in container.demux(stream): From 9de09dce5f7312fffcf2a25ae0f6bb9e7b60bc24 Mon Sep 17 00:00:00 2001 From: Santtu Keskinen Date: Tue, 9 Apr 2024 20:21:50 +0300 Subject: [PATCH 3/4] More linter changes --- av/bitstream.pxd | 1 + av/bitstream.pyx | 3 +-- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/av/bitstream.pxd b/av/bitstream.pxd index ae88568c1..2ebc1da75 100644 --- a/av/bitstream.pxd +++ b/av/bitstream.pxd @@ -2,6 +2,7 @@ cimport libav as lib from av.packet cimport Packet + cdef class BitStreamFilterContext: cdef const lib.AVBSFContext *ptr diff --git a/av/bitstream.pyx b/av/bitstream.pyx index af04f2e01..949d53320 100644 --- a/av/bitstream.pyx +++ b/av/bitstream.pyx @@ -1,6 +1,5 @@ -from libc.errno cimport EAGAIN - cimport libav as lib +from libc.errno cimport EAGAIN from av.error cimport err_check from av.packet cimport Packet From ba7cdd95df7bb497f49f491e065d3ae4e3248cc5 Mon Sep 17 00:00:00 2001 From: WyattBlue Date: Tue, 9 Apr 2024 16:14:56 -0400 Subject: [PATCH 4/4] Add type hints, docstring --- av/bitstream.pyi | 10 +++++++++ av/bitstream.pyx | 48 +++++++++++++++++++++++------------------ av/packet.pyi | 6 +++--- docs/api/bitstream.rst | 9 ++++++++ docs/index.rst | 1 + tests/test_bitstream.py | 18 +++++++++------- 6 files changed, 60 insertions(+), 32 deletions(-) create mode 100644 av/bitstream.pyi create mode 100644 docs/api/bitstream.rst 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 index 949d53320..b1fedb761 100644 --- a/av/bitstream.pyx +++ b/av/bitstream.pyx @@ -7,7 +7,11 @@ 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 @@ -15,6 +19,7 @@ cdef class BitStreamFilterContext: 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) @@ -22,6 +27,7 @@ cdef class BitStreamFilterContext: 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) @@ -30,36 +36,35 @@ cdef class BitStreamFilterContext: if self.ptr: lib.av_bsf_free(&self.ptr) - def _send(self, Packet packet=None): - cdef int res - with nogil: - res = lib.av_bsf_send_packet(self.ptr, packet.ptr if packet is not None else NULL) - err_check(res) - - def _recv(self): - cdef Packet packet = Packet() + 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_receive_packet(self.ptr, packet.ptr) - if res == -EAGAIN or res == lib.AVERROR_EOF: - return + res = lib.av_bsf_send_packet(self.ptr, packet.ptr if packet is not None else NULL) err_check(res) - if not res: - return packet - - cpdef filter(self, Packet packet=None): - self._send(packet) - output = [] while True: - packet = self._recv() - if packet: - output.append(packet) - else: + 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() @@ -71,6 +76,7 @@ cdef get_filter_names(): 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/tests/test_bitstream.py b/tests/test_bitstream.py index da184ffa8..ddc1bfde2 100644 --- a/tests/test_bitstream.py +++ b/tests/test_bitstream.py @@ -6,14 +6,13 @@ class TestBitStreamFilters(TestCase): - - def test_filters_availible(self): + def test_filters_availible(self) -> None: self.assertIn("h264_mp4toannexb", bitstream_filters_available) - def test_filter_chomp(self): + def test_filter_chomp(self) -> None: ctx = BitStreamFilterContext("chomp") - src_packets = [Packet(b"\x0012345\0\0\0"), None] + 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 = [] @@ -23,16 +22,19 @@ def test_filter_chomp(self): self.assertEqual(len(result_packets), 1) self.assertEqual(bytes(result_packets[0]), b"\x0012345") - def test_filter_setts(self): + 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 = [] + result_packets: list[Packet] = [] for p in src_packets: result_packets.extend(ctx.filter(p)) @@ -40,8 +42,8 @@ def test_filter_setts(self): self.assertEqual(result_packets[0].pts, 0) self.assertEqual(result_packets[1].pts, 1) - def test_filter_h264_mp4toannexb(self): - def is_annexb(packet): + 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"