From ee5c40e41ca8badf4a14c38794cb19647084f656 Mon Sep 17 00:00:00 2001 From: Dave Johansen Date: Wed, 11 Sep 2024 18:58:55 -0600 Subject: [PATCH] Allow `extradata` to be set for encoders --- av/codec/context.pyx | 3 --- av/container/output.pyx | 21 +++++++++++++++------ tests/test_codec_context.py | 6 +++--- 3 files changed, 18 insertions(+), 12 deletions(-) diff --git a/av/codec/context.pyx b/av/codec/context.pyx index 66b21f9a3..e2e3632ee 100644 --- a/av/codec/context.pyx +++ b/av/codec/context.pyx @@ -189,9 +189,6 @@ cdef class CodecContext: @extradata.setter def extradata(self, data): - if not self.is_decoder: - raise ValueError("Can only set extradata for decoders.") - if data is None: lib.av_freep(&self.ptr.extradata) self.ptr.extradata_size = 0 diff --git a/av/container/output.pyx b/av/container/output.pyx index 00f5d55b3..717df8fb3 100644 --- a/av/container/output.pyx +++ b/av/container/output.pyx @@ -93,10 +93,14 @@ cdef class OutputContainer(Container): # Now lets set some more sane video defaults elif codec.type == lib.AVMEDIA_TYPE_VIDEO: codec_context.pix_fmt = lib.AV_PIX_FMT_YUV420P - codec_context.width = 640 - codec_context.height = 480 - codec_context.bit_rate = 1024000 - codec_context.bit_rate_tolerance = 128000 + codec_context.width = kwargs.pop("width", 640) + codec_context.height = kwargs.pop("height", 480) + codec_context.bit_rate = kwargs.pop("bit_rate", 1024000) + codec_context.bit_rate_tolerance = kwargs.pop("bit_rate_tolerance", 128000) + try: + to_avrational(kwargs.pop("time_base"), &codec_context.time_base) + except KeyError: + pass to_avrational(rate or 24, &codec_context.framerate) stream.avg_frame_rate = codec_context.framerate @@ -105,9 +109,14 @@ cdef class OutputContainer(Container): # Some sane audio defaults elif codec.type == lib.AVMEDIA_TYPE_AUDIO: codec_context.sample_fmt = codec.sample_fmts[0] - codec_context.bit_rate = 128000 - codec_context.bit_rate_tolerance = 32000 + codec_context.bit_rate = kwargs.pop("bit_rate", 128000) + codec_context.bit_rate_tolerance = kwargs.pop("bit_rate_tolerance", 32000) + try: + to_avrational(kwargs.pop("time_base"), &codec_context.time_base) + except KeyError: + pass codec_context.sample_rate = rate or 48000 + stream.time_base = codec_context.time_base lib.av_channel_layout_default(&codec_context.ch_layout, 2) # Some formats want stream headers to be separate diff --git a/tests/test_codec_context.py b/tests/test_codec_context.py index 7ada318d9..634fe0280 100644 --- a/tests/test_codec_context.py +++ b/tests/test_codec_context.py @@ -105,9 +105,9 @@ def test_encoder_extradata(self): self.assertEqual(ctx.extradata, None) self.assertEqual(ctx.extradata_size, 0) - with self.assertRaises(ValueError) as cm: - ctx.extradata = b"123" - self.assertEqual(str(cm.exception), "Can only set extradata for decoders.") + ctx.extradata = b"123" + self.assertEqual(ctx.extradata, b"123") + self.assertEqual(ctx.extradata_size, 3) def test_encoder_pix_fmt(self): ctx = av.codec.Codec("h264", "w").create()