Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
126 changes: 83 additions & 43 deletions av/container/output.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
from cython.cimports.av.packet import Packet
from cython.cimports.av.stream import Stream, wrap_stream
from cython.cimports.av.utils import dict_to_avdict, to_avrational
from cython.cimports.libc.stdint import uint8_t
from cython.cimports.libc.stdint import int64_t, uint8_t
from cython.cimports.libc.string import memcpy, memset


Expand Down Expand Up @@ -87,6 +87,8 @@ def __cinit__(self, *args, **kwargs):
self._buffered_packets = []
with cython.nogil:
self.packet_ptr = lib.av_packet_alloc()
if self.packet_ptr == cython.NULL:
raise MemoryError("Could not allocate packet")

def __del__(self):
close_output(self)
Expand Down Expand Up @@ -144,12 +146,22 @@ def add_stream(
has_time_base: cython.bint = "time_base" in kwargs
if has_time_base:
to_avrational(kwargs.pop("time_base"), cython.address(c_time_base))

c_width: cython.int = 0
c_height: cython.int = 0
c_bit_rate: int64_t = 0
c_bit_rate_tolerance: cython.int = 0
if codec.type == lib.AVMEDIA_TYPE_VIDEO:
to_avrational(rate or 24, cython.address(c_framerate))
elif codec.type == lib.AVMEDIA_TYPE_AUDIO and not (
rate is None or type(rate) is int
):
raise TypeError("audio stream `rate` must be: int | None")
c_width = kwargs.pop("width", 640)
c_height = kwargs.pop("height", 480)
c_bit_rate = kwargs.pop("bit_rate", 0)
c_bit_rate_tolerance = kwargs.pop("bit_rate_tolerance", 128000)
elif codec.type == lib.AVMEDIA_TYPE_AUDIO:
if not (rate is None or type(rate) is int):
raise TypeError("audio stream `rate` must be: int | None")
c_bit_rate = kwargs.pop("bit_rate", 0)
c_bit_rate_tolerance = kwargs.pop("bit_rate_tolerance", 32000)

# Create new stream in the AVFormatContext, set AVCodecContext values.
ctx: cython.pointer[lib.AVCodecContext] = lib.avcodec_alloc_context3(codec)
Expand All @@ -166,10 +178,10 @@ def add_stream(
# Now let's set some more sane video defaults
if codec.type == lib.AVMEDIA_TYPE_VIDEO:
ctx.pix_fmt = lib.AV_PIX_FMT_YUV420P
ctx.width = kwargs.pop("width", 640)
ctx.height = kwargs.pop("height", 480)
ctx.bit_rate = kwargs.pop("bit_rate", 0)
ctx.bit_rate_tolerance = kwargs.pop("bit_rate_tolerance", 128000)
ctx.width = c_width
ctx.height = c_height
ctx.bit_rate = c_bit_rate
ctx.bit_rate_tolerance = c_bit_rate_tolerance
ctx.framerate = c_framerate

stream.avg_frame_rate = ctx.framerate
Expand All @@ -188,8 +200,8 @@ def add_stream(
)
if out:
ctx.sample_fmt = cython.cast(cython.pointer[lib.AVSampleFormat], out)[0]
ctx.bit_rate = kwargs.pop("bit_rate", 0)
ctx.bit_rate_tolerance = kwargs.pop("bit_rate_tolerance", 32000)
ctx.bit_rate = c_bit_rate
ctx.bit_rate_tolerance = c_bit_rate_tolerance
ctx.sample_rate = 48000 if rate is None else rate
stream.time_base = ctx.time_base
lib.av_channel_layout_default(cython.address(ctx.ch_layout), 2)
Expand All @@ -202,9 +214,13 @@ def add_stream(
#
# Subsequent changes to the codec context will be applied just before
# encoding starts in `start_encoding()`.
err_check(lib.avcodec_parameters_from_context(stream.codecpar, ctx))
try:
err_check(lib.avcodec_parameters_from_context(stream.codecpar, ctx))
except Exception:
lib.avcodec_free_context(cython.address(ctx))
raise

# Construct the user-land stream
# Construct the user-land stream, which takes ownership of ctx.
py_codec_context: CodecContext = wrap_codec_context(ctx, codec, hwaccel)
py_stream: Stream = wrap_stream(self, stream, py_codec_context)
self.streams.add_stream(py_stream)
Expand Down Expand Up @@ -266,10 +282,15 @@ def add_mux_stream(self, codec_name: str, rate=None, **kwargs) -> Stream:
)

c_rate: lib.AVRational
if rate is not None:
if codec_type == lib.AVMEDIA_TYPE_VIDEO:
c_width: cython.int = 0
c_height: cython.int = 0
if codec_type == lib.AVMEDIA_TYPE_VIDEO:
if rate is not None:
to_avrational(rate, cython.address(c_rate))
elif codec_type == lib.AVMEDIA_TYPE_AUDIO and type(rate) is not int:
c_width = kwargs.pop("width", 0)
c_height = kwargs.pop("height", 0)
elif codec_type == lib.AVMEDIA_TYPE_AUDIO:
if rate is not None and type(rate) is not int:
raise TypeError("audio stream `rate` must be: int | None")

# Create stream with no codec context.
Expand All @@ -283,8 +304,8 @@ def add_mux_stream(self, codec_name: str, rate=None, **kwargs) -> Stream:
stream.codecpar.codec_type = codec_type

if codec_type == lib.AVMEDIA_TYPE_VIDEO:
stream.codecpar.width = kwargs.pop("width", 0)
stream.codecpar.height = kwargs.pop("height", 0)
stream.codecpar.width = c_width
stream.codecpar.height = c_height
if rate is not None:
stream.avg_frame_rate = c_rate
elif codec_type == lib.AVMEDIA_TYPE_AUDIO and rate is not None:
Expand Down Expand Up @@ -336,31 +357,42 @@ def add_stream_from_template(
)

# Create new stream in the AVFormatContext, set AVCodecContext values.
stream: cython.pointer[lib.AVStream] = lib.avformat_new_stream(self.ptr, codec)
ctx: cython.pointer[lib.AVCodecContext] = lib.avcodec_alloc_context3(codec)
if ctx == cython.NULL:
raise MemoryError("Could not allocate codec context")
stream: cython.pointer[lib.AVStream] = lib.avformat_new_stream(self.ptr, codec)
if stream == cython.NULL:
lib.avcodec_free_context(cython.address(ctx))
raise MemoryError("Could not allocate stream")

err_check(lib.avcodec_parameters_to_context(ctx, template.ptr.codecpar))
# Reset the codec tag assuming we are remuxing.
ctx.codec_tag = 0
try:
err_check(lib.avcodec_parameters_to_context(ctx, template.ptr.codecpar))
# Reset the codec tag assuming we are remuxing.
ctx.codec_tag = 0

# Copy the template's stream time_base
stream.time_base = template.ptr.time_base
ctx.time_base = template.ptr.time_base
# Copy the template's stream time_base
stream.time_base = template.ptr.time_base
ctx.time_base = template.ptr.time_base

# Some formats want stream headers to be separate
if self.ptr.oformat.flags & lib.AVFMT_GLOBALHEADER:
ctx.flags |= lib.AV_CODEC_FLAG_GLOBAL_HEADER
# Some formats want stream headers to be separate
if self.ptr.oformat.flags & lib.AVFMT_GLOBALHEADER:
ctx.flags |= lib.AV_CODEC_FLAG_GLOBAL_HEADER

# Copy flags If we're creating a new codec object. This fixes some muxing issues.
# Overwriting `ctx.flags |= lib.AV_CODEC_FLAG_GLOBAL_HEADER` is intentional.
if not opaque:
ctx.flags = template.codec_context.flags
# Copy flags If we're creating a new codec object. This fixes some
# muxing issues. Overwriting the flag set just above is intentional.
if not opaque:
ctx.flags = template.codec_context.flags

# Initialize stream codec parameters to populate the codec type. Subsequent changes to
# the codec context will be applied just before encoding starts in `start_encoding()`.
err_check(lib.avcodec_parameters_from_context(stream.codecpar, ctx))
# Initialize stream codec parameters to populate the codec type.
# Subsequent changes to the codec context will be applied just
# before encoding starts in `start_encoding()`.
err_check(lib.avcodec_parameters_from_context(stream.codecpar, ctx))
except Exception:
# Nothing owns ctx until wrap_codec_context() below.
lib.avcodec_free_context(cython.address(ctx))
raise

# Construct the user-land stream
# Construct the user-land stream, which takes ownership of ctx.
py_codec_context: CodecContext = wrap_codec_context(ctx, codec, None)
py_codec_context._ctxflags |= 1 # _template_initialized = True
py_stream: Stream = wrap_stream(self, stream, py_codec_context)
Expand Down Expand Up @@ -482,24 +514,32 @@ def add_data_stream(self, codec_name=None, options: dict | None = None):
f"{self.format.name!r} format does not support {codec_name!r} codec"
)

# Create new stream in the AVFormatContext
stream: cython.pointer[lib.AVStream] = lib.avformat_new_stream(self.ptr, codec)
if stream == cython.NULL:
raise MemoryError("Could not allocate stream")

# Set up codec context and parameters
# The context first, so a failure here does not orphan a stream.
ctx: cython.pointer[lib.AVCodecContext] = cython.NULL
if codec != cython.NULL:
ctx = lib.avcodec_alloc_context3(codec)
if ctx == cython.NULL:
raise MemoryError("Could not allocate codec context")

# Create new stream in the AVFormatContext
stream: cython.pointer[lib.AVStream] = lib.avformat_new_stream(self.ptr, codec)
if stream == cython.NULL:
if ctx != cython.NULL:
lib.avcodec_free_context(cython.address(ctx))
raise MemoryError("Could not allocate stream")

if codec != cython.NULL:
# Some formats want stream headers to be separate
if self.ptr.oformat.flags & lib.AVFMT_GLOBALHEADER:
ctx.flags |= lib.AV_CODEC_FLAG_GLOBAL_HEADER

# Initialize stream codec parameters
err_check(lib.avcodec_parameters_from_context(stream.codecpar, ctx))
try:
err_check(lib.avcodec_parameters_from_context(stream.codecpar, ctx))
except Exception:
# Nothing owns ctx until wrap_codec_context() below.
lib.avcodec_free_context(cython.address(ctx))
raise
else:
# No codec available - set basic parameters for data stream
stream.codecpar.codec_type = lib.AVMEDIA_TYPE_DATA
Expand Down
Loading