add read me
This commit is contained in:
474
venv/lib/python3.12/site-packages/av/container/output.py
Normal file
474
venv/lib/python3.12/site-packages/av/container/output.py
Normal file
@@ -0,0 +1,474 @@
|
||||
import os
|
||||
from fractions import Fraction
|
||||
|
||||
import cython
|
||||
from cython.cimports import libav as lib
|
||||
from cython.cimports.av.codec.codec import Codec
|
||||
from cython.cimports.av.codec.context import CodecContext, wrap_codec_context
|
||||
from cython.cimports.av.container.streams import StreamContainer
|
||||
from cython.cimports.av.dictionary import _Dictionary
|
||||
from cython.cimports.av.error import err_check
|
||||
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 av.dictionary import Dictionary
|
||||
|
||||
|
||||
@cython.cfunc
|
||||
def close_output(self: OutputContainer):
|
||||
self.streams = StreamContainer()
|
||||
if self._started and not self._done:
|
||||
# We must only ever call av_write_trailer *once*, otherwise we get a
|
||||
# segmentation fault. Therefore no matter whether it succeeds or not
|
||||
# we must absolutely set self._done.
|
||||
try:
|
||||
self.err_check(lib.av_write_trailer(self.ptr))
|
||||
finally:
|
||||
if self.file is None and not (self.ptr.oformat.flags & lib.AVFMT_NOFILE):
|
||||
lib.avio_closep(cython.address(self.ptr.pb))
|
||||
self._done = True
|
||||
|
||||
|
||||
@cython.cclass
|
||||
class OutputContainer(Container):
|
||||
def __cinit__(self, *args, **kwargs):
|
||||
self.streams = StreamContainer()
|
||||
self.metadata = {}
|
||||
with cython.nogil:
|
||||
self.packet_ptr = lib.av_packet_alloc()
|
||||
|
||||
def __dealloc__(self):
|
||||
close_output(self)
|
||||
with cython.nogil:
|
||||
lib.av_packet_free(cython.address(self.packet_ptr))
|
||||
|
||||
def add_stream(self, codec_name, rate=None, options: dict | None = None, **kwargs):
|
||||
"""add_stream(codec_name, rate=None)
|
||||
|
||||
Creates a new stream from a codec name and returns it.
|
||||
Supports video, audio, and subtitle streams.
|
||||
|
||||
:param codec_name: The name of a codec.
|
||||
:type codec_name: str
|
||||
:param dict options: Stream options.
|
||||
:param \\**kwargs: Set attributes for the stream.
|
||||
:rtype: The new :class:`~av.stream.Stream`.
|
||||
|
||||
"""
|
||||
|
||||
codec_obj: Codec = Codec(codec_name, "w")
|
||||
codec: cython.pointer[cython.const[lib.AVCodec]] = codec_obj.ptr
|
||||
|
||||
# Assert that this format supports the requested codec.
|
||||
if not lib.avformat_query_codec(
|
||||
self.ptr.oformat, codec.id, lib.FF_COMPLIANCE_NORMAL
|
||||
):
|
||||
raise ValueError(
|
||||
f"{self.format.name!r} format does not support {codec_obj.name!r} codec"
|
||||
)
|
||||
|
||||
# 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)
|
||||
|
||||
# Now lets 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)
|
||||
try:
|
||||
to_avrational(kwargs.pop("time_base"), cython.address(ctx.time_base))
|
||||
except KeyError:
|
||||
pass
|
||||
to_avrational(rate or 24, cython.address(ctx.framerate))
|
||||
|
||||
stream.avg_frame_rate = ctx.framerate
|
||||
stream.time_base = ctx.time_base
|
||||
|
||||
# Some sane audio defaults
|
||||
elif codec.type == lib.AVMEDIA_TYPE_AUDIO:
|
||||
ctx.sample_fmt = codec.sample_fmts[0]
|
||||
ctx.bit_rate = kwargs.pop("bit_rate", 0)
|
||||
ctx.bit_rate_tolerance = kwargs.pop("bit_rate_tolerance", 32000)
|
||||
try:
|
||||
to_avrational(kwargs.pop("time_base"), cython.address(ctx.time_base))
|
||||
except KeyError:
|
||||
pass
|
||||
|
||||
if rate is None:
|
||||
ctx.sample_rate = 48000
|
||||
elif type(rate) is int:
|
||||
ctx.sample_rate = rate
|
||||
else:
|
||||
raise TypeError("audio stream `rate` must be: int | None")
|
||||
stream.time_base = ctx.time_base
|
||||
lib.av_channel_layout_default(cython.address(ctx.ch_layout), 2)
|
||||
|
||||
# Some formats want stream headers to be separate
|
||||
if self.ptr.oformat.flags & lib.AVFMT_GLOBALHEADER:
|
||||
ctx.flags |= lib.AV_CODEC_FLAG_GLOBAL_HEADER
|
||||
|
||||
# Initialise 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))
|
||||
|
||||
# Construct the user-land stream
|
||||
py_codec_context: CodecContext = wrap_codec_context(ctx, codec, None)
|
||||
py_stream: Stream = wrap_stream(self, stream, py_codec_context)
|
||||
self.streams.add_stream(py_stream)
|
||||
|
||||
if options:
|
||||
py_stream.options.update(options)
|
||||
|
||||
for k, v in kwargs.items():
|
||||
setattr(py_stream, k, v)
|
||||
|
||||
return py_stream
|
||||
|
||||
def add_stream_from_template(
|
||||
self, template: Stream, opaque: bool | None = None, **kwargs
|
||||
):
|
||||
"""
|
||||
Creates a new stream from a template. Supports video, audio, subtitle, data and attachment streams.
|
||||
|
||||
:param template: Copy codec from another :class:`~av.stream.Stream` instance.
|
||||
:param opaque: If True, copy opaque data from the template's codec context.
|
||||
:param \\**kwargs: Set attributes for the stream.
|
||||
:rtype: The new :class:`~av.stream.Stream`.
|
||||
"""
|
||||
if opaque is None:
|
||||
opaque = template.type != "video"
|
||||
|
||||
if template.codec_context is None:
|
||||
return self._add_stream_without_codec_from_template(template, **kwargs)
|
||||
|
||||
codec_obj: Codec
|
||||
if opaque: # Copy ctx from template.
|
||||
codec_obj = template.codec_context.codec
|
||||
else: # Construct new codec object.
|
||||
codec_obj = Codec(template.codec_context.codec.name, "w")
|
||||
|
||||
codec: cython.pointer[cython.const[lib.AVCodec]] = codec_obj.ptr
|
||||
|
||||
# Assert that this format supports the requested codec.
|
||||
if not lib.avformat_query_codec(
|
||||
self.ptr.oformat, codec.id, lib.FF_COMPLIANCE_NORMAL
|
||||
):
|
||||
raise ValueError(
|
||||
f"{self.format.name!r} format does not support {codec_obj.name!r} codec"
|
||||
)
|
||||
|
||||
# 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)
|
||||
|
||||
err_check(lib.avcodec_parameters_to_context(ctx, template.ptr.codecpar))
|
||||
# Reset the codec tag assuming we are remuxing.
|
||||
ctx.codec_tag = 0
|
||||
|
||||
# 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
|
||||
|
||||
# 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))
|
||||
|
||||
# Construct the user-land stream
|
||||
py_codec_context: CodecContext = wrap_codec_context(ctx, codec, None)
|
||||
py_stream: Stream = wrap_stream(self, stream, py_codec_context)
|
||||
self.streams.add_stream(py_stream)
|
||||
|
||||
for k, v in kwargs.items():
|
||||
setattr(py_stream, k, v)
|
||||
|
||||
return py_stream
|
||||
|
||||
def _add_stream_without_codec_from_template(
|
||||
self, template: Stream, **kwargs
|
||||
) -> Stream:
|
||||
codec_type: cython.int = template.ptr.codecpar.codec_type
|
||||
if codec_type not in {lib.AVMEDIA_TYPE_ATTACHMENT, lib.AVMEDIA_TYPE_DATA}:
|
||||
raise ValueError(
|
||||
f"template stream of type {template.type} has no codec context"
|
||||
)
|
||||
|
||||
stream: cython.pointer[lib.AVStream] = lib.avformat_new_stream(
|
||||
self.ptr, cython.NULL
|
||||
)
|
||||
if stream == cython.NULL:
|
||||
raise MemoryError("Could not allocate stream")
|
||||
|
||||
err_check(lib.avcodec_parameters_copy(stream.codecpar, template.ptr.codecpar))
|
||||
|
||||
# Mirror basic properties that are not derived from a codec context.
|
||||
stream.time_base = template.ptr.time_base
|
||||
stream.start_time = template.ptr.start_time
|
||||
stream.duration = template.ptr.duration
|
||||
stream.disposition = template.ptr.disposition
|
||||
|
||||
py_stream: Stream = wrap_stream(self, stream, None)
|
||||
self.streams.add_stream(py_stream)
|
||||
|
||||
py_stream.metadata = dict(template.metadata)
|
||||
|
||||
for k, v in kwargs.items():
|
||||
setattr(py_stream, k, v)
|
||||
|
||||
return py_stream
|
||||
|
||||
def add_attachment(self, name: str, mimetype: str, data: bytes):
|
||||
"""
|
||||
Create an attachment stream and embed its payload into the container header.
|
||||
|
||||
- Only supported by formats that support attachments (e.g. Matroska).
|
||||
- No per-packet muxing is required; attachments are written at header time.
|
||||
"""
|
||||
# Create stream with no codec (attachments are codec-less).
|
||||
stream: cython.pointer[lib.AVStream] = lib.avformat_new_stream(
|
||||
self.ptr, cython.NULL
|
||||
)
|
||||
if stream == cython.NULL:
|
||||
raise MemoryError("Could not allocate stream")
|
||||
|
||||
stream.codecpar.codec_type = lib.AVMEDIA_TYPE_ATTACHMENT
|
||||
stream.codecpar.codec_id = lib.AV_CODEC_ID_NONE
|
||||
|
||||
# Allocate and copy payload into codecpar.extradata.
|
||||
payload_size: cython.size_t = len(data)
|
||||
if payload_size:
|
||||
buf = cython.cast(cython.p_uchar, lib.av_malloc(payload_size + 1))
|
||||
if buf == cython.NULL:
|
||||
raise MemoryError("Could not allocate attachment data")
|
||||
# Copy bytes.
|
||||
for i in range(payload_size):
|
||||
buf[i] = data[i]
|
||||
buf[payload_size] = 0
|
||||
stream.codecpar.extradata = cython.cast(cython.p_uchar, buf)
|
||||
stream.codecpar.extradata_size = payload_size
|
||||
|
||||
# Wrap as user-land stream.
|
||||
meta_ptr = cython.address(stream.metadata)
|
||||
err_check(lib.av_dict_set(meta_ptr, b"filename", name.encode(), 0))
|
||||
mime_bytes = mimetype.encode()
|
||||
err_check(lib.av_dict_set(meta_ptr, b"mimetype", mime_bytes, 0))
|
||||
|
||||
py_stream: Stream = wrap_stream(self, stream, None)
|
||||
self.streams.add_stream(py_stream)
|
||||
return py_stream
|
||||
|
||||
def add_data_stream(self, codec_name=None, options: dict | None = None):
|
||||
"""add_data_stream(codec_name=None)
|
||||
|
||||
Creates a new data stream and returns it.
|
||||
|
||||
:param codec_name: Optional name of the data codec (e.g. 'klv')
|
||||
:type codec_name: str | None
|
||||
:param dict options: Stream options.
|
||||
:rtype: The new :class:`~av.data.stream.DataStream`.
|
||||
"""
|
||||
codec: cython.pointer[cython.const[lib.AVCodec]] = cython.NULL
|
||||
|
||||
if codec_name is not None:
|
||||
codec = lib.avcodec_find_encoder_by_name(codec_name.encode())
|
||||
if codec == cython.NULL:
|
||||
raise ValueError(f"Unknown data codec: {codec_name}")
|
||||
|
||||
# Assert that this format supports the requested codec
|
||||
if not lib.avformat_query_codec(
|
||||
self.ptr.oformat, codec.id, lib.FF_COMPLIANCE_NORMAL
|
||||
):
|
||||
raise ValueError(
|
||||
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 if we have a codec
|
||||
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")
|
||||
|
||||
# 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))
|
||||
else:
|
||||
# For raw data streams, just set the codec type
|
||||
stream.codecpar.codec_type = lib.AVMEDIA_TYPE_DATA
|
||||
|
||||
# Construct the user-land stream
|
||||
py_codec_context: CodecContext | None = None
|
||||
if ctx != cython.NULL:
|
||||
py_codec_context = wrap_codec_context(ctx, codec, None)
|
||||
|
||||
py_stream: Stream = wrap_stream(self, stream, py_codec_context)
|
||||
self.streams.add_stream(py_stream)
|
||||
|
||||
if options:
|
||||
py_stream.options.update(options)
|
||||
|
||||
return py_stream
|
||||
|
||||
@cython.ccall
|
||||
def start_encoding(self):
|
||||
"""Write the file header! Called automatically."""
|
||||
if self._started:
|
||||
return
|
||||
|
||||
# TODO: This does NOT handle options coming from 3 sources.
|
||||
# This is only a rough approximation of what would be cool to do.
|
||||
used_options: set = set()
|
||||
stream: Stream
|
||||
|
||||
# Finalize and open all streams.
|
||||
for stream in self.streams:
|
||||
ctx = stream.codec_context
|
||||
# Skip codec context handling for streams without codecs (e.g. data/attachments).
|
||||
if ctx is None:
|
||||
if stream.type not in {"data", "attachment"}:
|
||||
raise ValueError(f"Stream {stream.index} has no codec context")
|
||||
else:
|
||||
if not ctx.is_open:
|
||||
for k, v in self.options.items():
|
||||
ctx.options.setdefault(k, v)
|
||||
ctx.open()
|
||||
|
||||
# Track option consumption.
|
||||
for k in self.options:
|
||||
if k not in ctx.options:
|
||||
used_options.add(k)
|
||||
|
||||
stream._finalize_for_output()
|
||||
|
||||
# Open the output file, if needed.
|
||||
name_obj: bytes = os.fsencode(self.name if self.file is None else "")
|
||||
name: cython.p_char = name_obj
|
||||
if self.ptr.pb == cython.NULL and not self.ptr.oformat.flags & lib.AVFMT_NOFILE:
|
||||
err_check(
|
||||
lib.avio_open(cython.address(self.ptr.pb), name, lib.AVIO_FLAG_WRITE)
|
||||
)
|
||||
|
||||
# Copy the metadata dict.
|
||||
dict_to_avdict(
|
||||
cython.address(self.ptr.metadata),
|
||||
self.metadata,
|
||||
encoding=self.metadata_encoding,
|
||||
errors=self.metadata_errors,
|
||||
)
|
||||
|
||||
all_options: _Dictionary = Dictionary(self.options, self.container_options)
|
||||
options: _Dictionary = all_options.copy()
|
||||
self.err_check(lib.avformat_write_header(self.ptr, cython.address(options.ptr)))
|
||||
|
||||
# Track option usage...
|
||||
for k in all_options:
|
||||
if k not in options:
|
||||
used_options.add(k)
|
||||
|
||||
# ... and warn if any weren't used.
|
||||
unused_options = {
|
||||
k: v for k, v in self.options.items() if k not in used_options
|
||||
}
|
||||
if unused_options:
|
||||
import logging
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
log.warning("Some options were not used: %s" % unused_options)
|
||||
|
||||
self._started = True
|
||||
|
||||
@property
|
||||
def supported_codecs(self):
|
||||
"""
|
||||
Returns a set of all codecs this format supports.
|
||||
"""
|
||||
result: set = set()
|
||||
codec: cython.pointer[cython.const[lib.AVCodec]] = cython.NULL
|
||||
opaque: cython.p_void = cython.NULL
|
||||
|
||||
while True:
|
||||
codec = lib.av_codec_iterate(cython.address(opaque))
|
||||
if codec == cython.NULL:
|
||||
break
|
||||
|
||||
if (
|
||||
lib.avformat_query_codec(
|
||||
self.ptr.oformat, codec.id, lib.FF_COMPLIANCE_NORMAL
|
||||
)
|
||||
== 1
|
||||
):
|
||||
result.add(codec.name)
|
||||
|
||||
return result
|
||||
|
||||
@property
|
||||
def default_video_codec(self):
|
||||
"""
|
||||
Returns the default video codec this container recommends.
|
||||
"""
|
||||
return lib.avcodec_get_name(self.format.optr.video_codec)
|
||||
|
||||
@property
|
||||
def default_audio_codec(self):
|
||||
"""
|
||||
Returns the default audio codec this container recommends.
|
||||
"""
|
||||
return lib.avcodec_get_name(self.format.optr.audio_codec)
|
||||
|
||||
@property
|
||||
def default_subtitle_codec(self):
|
||||
"""
|
||||
Returns the default subtitle codec this container recommends.
|
||||
"""
|
||||
return lib.avcodec_get_name(self.format.optr.subtitle_codec)
|
||||
|
||||
def close(self):
|
||||
close_output(self)
|
||||
|
||||
def mux(self, packets):
|
||||
# We accept either a Packet, or a sequence of packets. This should smooth out
|
||||
# the transition to the new encode API which returns a sequence of packets.
|
||||
if isinstance(packets, Packet):
|
||||
self.mux_one(packets)
|
||||
else:
|
||||
for packet in packets:
|
||||
self.mux_one(packet)
|
||||
|
||||
def mux_one(self, packet: Packet):
|
||||
self.start_encoding()
|
||||
|
||||
# Assert the packet is in stream time.
|
||||
if (
|
||||
packet.ptr.stream_index < 0
|
||||
or cython.cast(cython.uint, packet.ptr.stream_index) >= self.ptr.nb_streams
|
||||
):
|
||||
raise ValueError("Bad Packet stream_index.")
|
||||
|
||||
stream: cython.pointer[lib.AVStream] = self.ptr.streams[packet.ptr.stream_index]
|
||||
packet._rebase_time(stream.time_base)
|
||||
|
||||
# Make another reference to the packet, as `av_interleaved_write_frame()`
|
||||
# takes ownership of the reference.
|
||||
self.err_check(lib.av_packet_ref(self.packet_ptr, packet.ptr))
|
||||
|
||||
with cython.nogil:
|
||||
ret: cython.int = lib.av_interleaved_write_frame(self.ptr, self.packet_ptr)
|
||||
self.err_check(ret)
|
||||
Reference in New Issue
Block a user