diff --git a/CHANGELOG.rst b/CHANGELOG.rst index e02cb1dd8..b5f2811cd 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -37,6 +37,7 @@ Major: Features: +- Add ``OutputContainer.add_mux_stream()`` for creating codec-context-free streams, enabling muxing of pre-encoded packets without an encoder, addressing :issue:`1970` by :gh-user:`WyattBlue`. - Use zero-copy for Packet init from buffer data by :gh-user:`WyattBlue` in (:pr:`2199`). - Expose AVIndexEntry by :gh-user:`Queuecumber` in (:pr:`2136`). - Preserving hardware memory during cuvid decoding, exporting/importing via dlpack by :gh-user:`WyattBlue` in (:pr:`2155`). diff --git a/av/audio/stream.py b/av/audio/stream.py index 7c150c84b..e32f7f759 100644 --- a/av/audio/stream.py +++ b/av/audio/stream.py @@ -6,6 +6,8 @@ @cython.cclass class AudioStream(Stream): def __repr__(self): + if self.codec_context is None: + return f" at 0x{id(self):x}>" form = self.format.name if self.format else None return ( f" Stream: + """add_mux_stream(codec_name, rate=None) + + Creates a new stream for muxing pre-encoded data without creating a + :class:`.CodecContext`. Use this when you want to mux packets that were + already encoded externally and no encoding/decoding is needed. + + :param codec_name: The name of a codec. + :type codec_name: str + :param \\**kwargs: Set attributes for the stream (e.g. ``width``, ``height``, + ``time_base``). + :rtype: The new :class:`~av.stream.Stream`. + + """ + # Find the codec to get its id and type (try encoder first, then decoder). + codec_name_bytes: bytes = codec_name.encode() + codec: cython.pointer[cython.const[lib.AVCodec]] = ( + lib.avcodec_find_encoder_by_name(codec_name_bytes) + ) + codec_descriptor: cython.pointer[cython.const[lib.AVCodecDescriptor]] = ( + cython.NULL + ) + if codec == cython.NULL: + codec = lib.avcodec_find_decoder_by_name(codec_name_bytes) + if codec == cython.NULL: + codec_descriptor = lib.avcodec_descriptor_get_by_name(codec_name_bytes) + if codec_descriptor == cython.NULL: + raise ValueError(f"Unknown codec: {codec_name!r}") + + codec_id: lib.AVCodecID + codec_type: lib.AVMediaType + if codec != cython.NULL: + codec_id = codec.id + codec_type = codec.type + else: + codec_id = codec_descriptor.id + codec_type = codec_descriptor.type + + # 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 stream with 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") + + stream.codecpar.codec_id = codec_id + 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) + if rate is not None: + to_avrational(rate, cython.address(stream.avg_frame_rate)) + elif codec_type == lib.AVMEDIA_TYPE_AUDIO: + if rate is not None: + if type(rate) is int: + stream.codecpar.sample_rate = rate + else: + raise TypeError("audio stream `rate` must be: int | None") + + # Construct the user-land stream (no codec context). + py_stream: Stream = wrap_stream(self, stream, None) + self.streams.add_stream(py_stream) + + 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 ): @@ -291,13 +368,12 @@ def add_data_stream(self, codec_name=None, options: dict | None = None): ) if codec_name is not None: - codec = lib.avcodec_find_encoder_by_name(codec_name.encode()) + codec_name_bytes: bytes = codec_name.encode() + codec = lib.avcodec_find_encoder_by_name(codec_name_bytes) if codec == cython.NULL: - codec = lib.avcodec_find_decoder_by_name(codec_name.encode()) + codec = lib.avcodec_find_decoder_by_name(codec_name_bytes) if codec == cython.NULL: - codec_descriptor = lib.avcodec_descriptor_get_by_name( - codec_name.encode() - ) + codec_descriptor = lib.avcodec_descriptor_get_by_name(codec_name_bytes) if codec_descriptor == cython.NULL: raise ValueError(f"Unknown data codec: {codec_name}") @@ -361,22 +437,17 @@ def start_encoding(self): # 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) + if ctx is not None and not ctx.is_open: + for k, v in self.options.items(): + ctx.options.setdefault(k, v) - if not ctx._template_initialized: - ctx.open() + if not ctx._template_initialized: + ctx.open() - # Track option consumption. - for k in self.options: - if k not in ctx.options: - used_options.add(k) + # Track option consumption. + for k in self.options: + if k not in ctx.options: + used_options.add(k) stream._finalize_for_output() diff --git a/av/container/output.pyi b/av/container/output.pyi index 3fa18243c..de2e4950a 100644 --- a/av/container/output.pyi +++ b/av/container/output.pyi @@ -39,6 +39,12 @@ class OutputContainer(Container): options: dict[str, str] | None = None, **kwargs, ) -> VideoStream | AudioStream | SubtitleStream: ... + def add_mux_stream( + self, + codec_name: str, + rate: Fraction | int | None = None, + **kwargs, + ) -> Stream: ... def add_stream_from_template( self, template: _StreamT, opaque: bool | None = None, **kwargs ) -> _StreamT: ... diff --git a/av/video/stream.py b/av/video/stream.py index 658df3d5c..d5b2c106d 100644 --- a/av/video/stream.py +++ b/av/video/stream.py @@ -8,6 +8,8 @@ @cython.cclass class VideoStream(Stream): def __repr__(self): + if self.codec_context is None: + return f" at 0x{id(self):x}>" return ( f" None: packet_count += 1 assert packet_count > 50 + + +def test_add_mux_stream_video() -> None: + """add_mux_stream creates a video stream without a CodecContext.""" + input_path = av.datasets.curated("pexels/time-lapse-video-of-night-sky-857195.mp4") + + buf = io.BytesIO() + with av.open(input_path) as input_: + in_stream = input_.streams.video[0] + width = in_stream.codec_context.width + height = in_stream.codec_context.height + + with av.open(buf, "w", format="mp4") as output: + out_stream = output.add_mux_stream( + in_stream.codec_context.name, width=width, height=height + ) + assert out_stream.codec_context is None + assert out_stream.type == "video" + + out_stream.time_base = in_stream.time_base + + for packet in input_.demux(in_stream): + if packet.dts is None: + continue + packet.stream = out_stream + output.mux(packet) + + buf.seek(0) + with av.open(buf) as result: + assert len(result.streams.video) == 1 + assert result.streams.video[0].codec_context.width == width + assert result.streams.video[0].codec_context.height == height + + +def test_add_mux_stream_no_codec_context() -> None: + """add_mux_stream streams have no codec context and repr does not crash.""" + buf = io.BytesIO() + with av.open(buf, "w", format="mp4") as output: + video_stream = output.add_mux_stream("h264", width=1920, height=1080) + audio_stream = output.add_mux_stream("aac", rate=44100) + + assert video_stream.codec_context is None + assert audio_stream.codec_context is None + # repr should not crash + assert "video/" in repr(video_stream) + assert "audio/" in repr(audio_stream)