diff --git a/benchmarks/bench_blocks.py b/benchmarks/bench_blocks.py index a11a202b2..01ea29596 100644 --- a/benchmarks/bench_blocks.py +++ b/benchmarks/bench_blocks.py @@ -15,7 +15,7 @@ import torch from torchcodec.decoders import VideoDecoder -from torchcodec.decoders._blocks import ColorConverter, PacketDecoder, VideoDemuxer +from torchcodec.decoders._blocks import ColorConverter, VideoDemuxer, VideoPacketDecoder # Kept minimal on purpose; the filename is derived from exactly these. _DURATION_S = 10 @@ -113,7 +113,7 @@ def _consume(frames): def _decode_sequential(path, device="cpu"): demuxer = VideoDemuxer(path) - decoder = PacketDecoder(demuxer, device=device) + decoder = VideoPacketDecoder(demuxer, device=device) converter = ColorConverter(device=device) _consume(_convert(converter, _decode(decoder, _demux(demuxer)))) @@ -121,7 +121,7 @@ def _decode_sequential(path, device="cpu"): def _decode_prefetch_frames(path, device="cpu"): # [demux + decode] on one thread || [color-convert] on another. demuxer = VideoDemuxer(path) - decoder = PacketDecoder(demuxer, device=device) + decoder = VideoPacketDecoder(demuxer, device=device) converter = ColorConverter(device=device) frames = prefetch(_decode(decoder, _demux(demuxer))) _consume(_convert(converter, frames)) @@ -130,7 +130,7 @@ def _decode_prefetch_frames(path, device="cpu"): def _decode_prefetch_packets(path, device="cpu"): # [demux] on one thread || [decode + color-convert] on another. demuxer = VideoDemuxer(path) - decoder = PacketDecoder(demuxer, device=device) + decoder = VideoPacketDecoder(demuxer, device=device) converter = ColorConverter(device=device) packets = prefetch(_demux(demuxer)) _consume(_convert(converter, _decode(decoder, packets))) @@ -139,7 +139,7 @@ def _decode_prefetch_packets(path, device="cpu"): def _decode_prefetch_packets_and_frames(path, device="cpu"): # [demux] || [decode] || [color-convert], each on its own thread. demuxer = VideoDemuxer(path) - decoder = PacketDecoder(demuxer, device=device) + decoder = VideoPacketDecoder(demuxer, device=device) converter = ColorConverter(device=device) packets = prefetch(_demux(demuxer)) frames = prefetch(_decode(decoder, packets)) diff --git a/examples/decoding/blocks.py b/examples/decoding/blocks.py index f29701444..5c19a382c 100644 --- a/examples/decoding/blocks.py +++ b/examples/decoding/blocks.py @@ -21,7 +21,7 @@ .. code-block:: - VideoDemuxer -> PacketDecoder -> ColorConverter + VideoDemuxer -> VideoPacketDecoder -> ColorConverter Packet RawFrame RGB Frame The blocks are passive: they never create threads, and they release the GIL. @@ -63,14 +63,14 @@ # it can output a frame, and it buffers a few frames that ``drain()`` returns # at the end. # -# ``PacketDecoder`` and ``ColorConverter`` both accept ``device="cuda"``: +# ``VideoPacketDecoder`` and ``ColorConverter`` both accept ``device="cuda"``: # decoding then runs on NVDEC and the color conversion on the GPU, and the # frames never leave the device. Demuxing always happens on the CPU. Left # unspecified, ``device`` is the current default device. -from torchcodec.decoders._blocks import ColorConverter, VideoDemuxer, PacketDecoder +from torchcodec.decoders._blocks import ColorConverter, VideoDemuxer, VideoPacketDecoder demuxer = VideoDemuxer(video_path) -packet_decoder = PacketDecoder(demuxer, device=device) +packet_decoder = VideoPacketDecoder(demuxer, device=device) color_converter = ColorConverter(device=device) frames = [] @@ -134,7 +134,7 @@ def drain(): def sequential(): # demux -> decode -> color-convert, all on the calling thread. demuxer = VideoDemuxer(video_path) - packet_decoder = PacketDecoder(demuxer, device=device) + packet_decoder = VideoPacketDecoder(demuxer, device=device) color_converter = ColorConverter(device=device) return color_convert(color_converter, decode(packet_decoder, demux(demuxer))) @@ -142,7 +142,7 @@ def sequential(): def convert_on_own_thread(): # [demux + decode] on one thread || [color-convert] on another. demuxer = VideoDemuxer(video_path) - packet_decoder = PacketDecoder(demuxer, device=device) + packet_decoder = VideoPacketDecoder(demuxer, device=device) color_converter = ColorConverter(device=device) raw_frames = prefetch(decode(packet_decoder, demux(demuxer))) return color_convert(color_converter, raw_frames) @@ -153,7 +153,7 @@ def demux_on_own_thread(): # natural split on CUDA: demuxing is CPU and I/O work, while decoding and # color conversion both happen on the GPU, so they belong together. demuxer = VideoDemuxer(video_path) - packet_decoder = PacketDecoder(demuxer, device=device) + packet_decoder = VideoPacketDecoder(demuxer, device=device) color_converter = ColorConverter(device=device) packets = prefetch(demux(demuxer)) return color_convert(color_converter, decode(packet_decoder, packets)) @@ -179,9 +179,9 @@ def demux_on_own_thread(): # drop them until you reach the timestamp you asked for. # # The seek also invalidates the frames the decoder is holding on to, so the -# ``PacketDecoder`` must be ``reset()``. +# ``VideoPacketDecoder`` must be ``reset()``. demuxer = VideoDemuxer(video_path) -packet_decoder = PacketDecoder(demuxer, device=device) +packet_decoder = VideoPacketDecoder(demuxer, device=device) color_converter = ColorConverter(device=device) seconds = 2.5 @@ -212,7 +212,7 @@ def demux_on_own_thread(): # costs one pass over the file, and it leaves the demuxer back at the start. demuxer = VideoDemuxer(video_path) index = demuxer.scan() -packet_decoder = PacketDecoder(demuxer, device=device) +packet_decoder = VideoPacketDecoder(demuxer, device=device) color_converter = ColorConverter(device=device) print(f"{len(index)} frames at {index.average_fps} fps, " @@ -281,7 +281,7 @@ def demux_on_own_thread(): # Color conversion is optional. A ``RawFrame`` can hand out the decoder's own # planes as tensor views, with no copy and no conversion. demuxer = VideoDemuxer(video_path) -packet_decoder = PacketDecoder(demuxer, device=device) +packet_decoder = VideoPacketDecoder(demuxer, device=device) raw_frame = next(decode(packet_decoder, demux(demuxer))) Y, U, V = raw_frame.planes @@ -353,7 +353,7 @@ def upsample(plane): ) hdr_demuxer = VideoDemuxer(hdr_video_path) -hdr_packet_decoder = PacketDecoder(hdr_demuxer, device=device) +hdr_packet_decoder = VideoPacketDecoder(hdr_demuxer, device=device) hdr_raw = next(decode(hdr_packet_decoder, demux(hdr_demuxer))) hdr_Y = hdr_raw.planes[0] @@ -413,7 +413,7 @@ def start_live_stream(): # The blocks just stream it, and we stop whenever we want: ffmpeg = start_live_stream() demuxer = VideoDemuxer(fifo_path) -packet_decoder = PacketDecoder(demuxer, device=device) +packet_decoder = VideoPacketDecoder(demuxer, device=device) color_converter = ColorConverter(device=device) frames = [] diff --git a/src/torchcodec/_core/AudioCommon.cpp b/src/torchcodec/_core/AudioCommon.cpp new file mode 100644 index 000000000..64abd9b55 --- /dev/null +++ b/src/torchcodec/_core/AudioCommon.cpp @@ -0,0 +1,38 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include "AudioCommon.h" + +namespace facebook::torchcodec { + +torch::headeronly::ScalarType sample_format_dtype( + AVSampleFormat sample_format) { + switch (av_get_packed_sample_fmt(sample_format)) { + case AV_SAMPLE_FMT_U8: + return kStableUInt8; + case AV_SAMPLE_FMT_S16: + return kStableInt16; + case AV_SAMPLE_FMT_S32: + return kStableInt32; + case AV_SAMPLE_FMT_S64: + return kStableInt64; + case AV_SAMPLE_FMT_FLT: + return kStableFloat32; + case AV_SAMPLE_FMT_DBL: + return kStableFloat64; + default: + break; + } + const char* name = av_get_sample_fmt_name(sample_format); + STD_TORCH_CHECK( + false, + "Unsupported sample format '", + name == nullptr ? "unknown" : name, + "'."); + return kStableUInt8; +} + +} // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/AudioCommon.h b/src/torchcodec/_core/AudioCommon.h new file mode 100644 index 000000000..482bd86b4 --- /dev/null +++ b/src/torchcodec/_core/AudioCommon.h @@ -0,0 +1,23 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#pragma once + +#include "FFMPEGCommon.h" +#include "StableABICompat.h" + +// Where FFmpeg's audio samples meet torch tensors. FFMPEGCommon deliberately +// knows nothing about tensors, and these helpers are shared by the decode, the +// conversion and the SingleStreamDecoder paths, so they live on their own. + +namespace facebook::torchcodec { + +// The dtype that holds `sample_format`'s samples exactly. Planar and packed +// variants of a format share a sample type, which is why this doesn't care +// which one it is given. +torch::headeronly::ScalarType sample_format_dtype(AVSampleFormat sample_format); + +} // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/PacketDecoder.cpp b/src/torchcodec/_core/PacketDecoder.cpp index 1f64165f9..2cf8bc800 100644 --- a/src/torchcodec/_core/PacketDecoder.cpp +++ b/src/torchcodec/_core/PacketDecoder.cpp @@ -6,7 +6,10 @@ #include "PacketDecoder.h" +#include "AudioCommon.h" + #include +#include namespace facebook::torchcodec { @@ -57,7 +60,13 @@ const AVCodec* find_decoder( PacketDecoder::PacketDecoder( const Demuxer& demuxer, const StableDevice& device, - std::optional ffmpeg_thread_count) { + std::optional ffmpeg_thread_count) + : media_type_(demuxer.media_type()) { + bool is_audio = media_type_ == AVMEDIA_TYPE_AUDIO; + STD_TORCH_CHECK( + !is_audio || device.type() == kStableCPU, + "Audio can only be decoded on the CPU."); + device_interface_ = create_device_interface(device); STD_TORCH_CHECK( device_interface_ != nullptr, @@ -65,18 +74,34 @@ PacketDecoder::PacketDecoder( AVStream* stream = demuxer.active_stream(); time_base_ = stream->time_base; + is_mpeg_ps_ = std::string_view(demuxer.format_context()->iformat->name) == "mpeg"; - if (const int32_t* matrix = get_display_matrix_from_stream(stream)) { + + if (is_audio) { + // Audio codecs are hardcoded to a single FFmpeg thread, see + // https://github.com/pytorch/torchcodec/issues/1253. + ffmpeg_thread_count = 1; + } else if (const int32_t* matrix = get_display_matrix_from_stream(stream)) { display_matrix_.emplace(); std::copy( matrix, matrix + display_matrix_->size(), display_matrix_->begin()); } + const AVCodec* av_codec = find_decoder(stream, device_interface_.get()); codec_context_ = create_and_open_codec_context( stream, av_codec, device_interface_.get(), ffmpeg_thread_count); device_interface_->initialize(codec_context_); + if (is_audio) { + // Nothing else to set up: unlike video, we hand out the samples in the + // codec's own format, so no conversion state is needed here. Note we + // deliberately do NOT set request_sample_fmt: what SingleStreamDecoder + // asks for (FLTP) is an optimization for its own conversion, and here it + // would hide what the codec natively produces. + return; + } + const AVPixFmtDescriptor* stream_desc = av_pix_fmt_desc_get(codec_context_->pix_fmt); int stream_bit_depth = stream_desc ? stream_desc->comp[0].depth : 8; @@ -133,10 +158,12 @@ int PacketDecoder::receive_frame(UniqueAVFrame& av_frame) { int status = device_interface_->receive_frame(av_frame); if (status == AVSUCCESS) { device_interface_->make_frame_standalone(av_frame); - // Attach a copy of the display matrix to the frame, so the ColorConverter - // can use it. - set_display_matrix_on_frame( - *av_frame, display_matrix_ ? display_matrix_->data() : nullptr); + if (media_type_ == AVMEDIA_TYPE_VIDEO) { + // Attach a copy of the display matrix to the frame, so the ColorConverter + // can use it. + set_display_matrix_on_frame( + *av_frame, display_matrix_ ? display_matrix_->data() : nullptr); + } } return status; } @@ -240,4 +267,73 @@ std::vector frame_planes( return planes; } +namespace { +// Scatters `num_channels`-interleaved samples into one contiguous row per +// channel. Templated on an integer of the right width rather than the actual +// sample type: we're only moving bytes around, so all that matters is size. +template +void deinterleave( + const uint8_t* src, + uint8_t* dst, + int num_channels, + int num_samples) { + const T* in = reinterpret_cast(src); + T* out = reinterpret_cast(dst); + for (int channel = 0; channel < num_channels; ++channel) { + T* row = out + static_cast(channel) * num_samples; + for (int sample = 0; sample < num_samples; ++sample) { + row[sample] = in[static_cast(sample) * num_channels + channel]; + } + } +} +} // namespace + +torch::stable::Tensor audio_samples(const AVFrame& av_frame) { + auto sample_format = static_cast(av_frame.format); + int num_channels = get_num_channels(av_frame); + int64_t num_samples = av_frame.nb_samples; + + torch::stable::Tensor samples = torch::stable::empty( + {num_channels, num_samples}, sample_format_dtype(sample_format)); + if (num_samples == 0) { + return samples; + } + + int bytes_per_sample = av_get_bytes_per_sample(sample_format); + auto* dst = static_cast(samples.mutable_data_ptr()); + int64_t bytes_per_channel = num_samples * bytes_per_sample; + + if (av_sample_fmt_is_planar(sample_format)) { + for (int channel = 0; channel < num_channels; ++channel) { + // extended_data rather than data: the latter only holds + // AV_NUM_DATA_POINTERS (8) pointers, and we support more channels. + std::memcpy( + dst + channel * bytes_per_channel, + av_frame.extended_data[channel], + bytes_per_channel); + } + } else { + const uint8_t* src = av_frame.extended_data[0]; + int num_samples_int = static_cast(num_samples); + switch (bytes_per_sample) { + case 1: + deinterleave(src, dst, num_channels, num_samples_int); + break; + case 2: + deinterleave(src, dst, num_channels, num_samples_int); + break; + case 4: + deinterleave(src, dst, num_channels, num_samples_int); + break; + case 8: + deinterleave(src, dst, num_channels, num_samples_int); + break; + default: + STD_TORCH_CHECK( + false, "Unexpected sample width: ", bytes_per_sample, " bytes."); + } + } + return samples; +} + } // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/PacketDecoder.h b/src/torchcodec/_core/PacketDecoder.h index 6ad481588..71d50e9ea 100644 --- a/src/torchcodec/_core/PacketDecoder.h +++ b/src/torchcodec/_core/PacketDecoder.h @@ -29,8 +29,9 @@ SharedAVCodecContext create_and_open_codec_context( DeviceInterface* device_interface, std::optional thread_count); -// Decode building block: turns compressed packets into decoded (YUV) frames. -// Configured from a Demuxer's active stream; stateful. Not thread-safe. +// Decode building block: turns compressed packets into decoded frames - (YUV) +// pictures for a video stream, samples in the codec's own format for an audio +// one. Configured from a Demuxer's active stream; stateful. Not thread-safe. class FORCE_PUBLIC_VISIBILITY PacketDecoder { public: explicit PacketDecoder( @@ -67,10 +68,15 @@ class FORCE_PUBLIC_VISIBILITY PacketDecoder { return time_base_; } + AVMediaType media_type() const { + return media_type_; + } + private: std::unique_ptr device_interface_; SharedAVCodecContext codec_context_; AVRational time_base_ = {}; + AVMediaType media_type_ = AVMEDIA_TYPE_VIDEO; // Stamped onto every frame we hand out, so downstream blocks can read the // rotation off the frame itself instead of knowing about the stream. Held by // value: we're only handed the Demuxer at construction and it may well be @@ -97,6 +103,8 @@ struct FrameMetadata { double rotation_degrees = 0; }; +// TODO_API_BREAKDOWN CC P1 these should bet get_* + // Describes `av_frame` without touching its samples. Unlike frame_planes(), // this works for every pixel format, so callers can ask what a frame is before // asking for views they may not be able to get. @@ -111,4 +119,16 @@ FORCE_PUBLIC_VISIBILITY std::vector frame_planes( const StableDevice& device, const torch::stable::Tensor& tensor_handle); +// A decoded audio frame's samples as a contiguous [num_channels, num_samples] +// tensor whose dtype is the frame's own sample type: uint8 for u8, int16 for +// s16, float32 for flt, and so on, planar or not. This is a copy rather than a +// view: planar formats put each channel in its own allocation and packed ones +// interleave them, so neither is a [C, N] tensor as it stands. An audio frame +// is a few kB, so normalizing here buys a uniform layout for the price of a +// memcpy - and it means a converter can treat the result as planar-of-dtype. +// TODO_API_BREAKDOWN DESIGN P1: do we want to copy? Should we just keep the +// original layout? +FORCE_PUBLIC_VISIBILITY torch::stable::Tensor audio_samples( + const AVFrame& av_frame); + } // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/StableABICompat.h b/src/torchcodec/_core/StableABICompat.h index 514482f77..fe1a4ccc3 100644 --- a/src/torchcodec/_core/StableABICompat.h +++ b/src/torchcodec/_core/StableABICompat.h @@ -63,6 +63,7 @@ constexpr auto kStableXPU = torch::headeronly::DeviceType::XPU; // Scalar type constants constexpr auto kStableUInt8 = torch::headeronly::ScalarType::Byte; constexpr auto kStableUInt16 = torch::headeronly::ScalarType::UInt16; +constexpr auto kStableInt16 = torch::headeronly::ScalarType::Short; constexpr auto kStableInt32 = torch::headeronly::ScalarType::Int; constexpr auto kStableInt64 = torch::headeronly::ScalarType::Long; constexpr auto kStableFloat32 = torch::headeronly::ScalarType::Float; diff --git a/src/torchcodec/_core/_ffmpeg_op_names.py b/src/torchcodec/_core/_ffmpeg_op_names.py index 8523199c9..194a9803a 100644 --- a/src/torchcodec/_core/_ffmpeg_op_names.py +++ b/src/torchcodec/_core/_ffmpeg_op_names.py @@ -41,6 +41,7 @@ "_blocks_packet_decoder_send_eof", "_blocks_packet_decoder_reset", "_blocks_packet_decoder_receive_frame", + "_blocks_audio_packet_decoder_receive_frame", "_blocks_create_color_converter", "_blocks_convert_frame", "_blocks_frame_metadata", diff --git a/src/torchcodec/_core/_ffmpeg_ops.py b/src/torchcodec/_core/_ffmpeg_ops.py index d22a3b8f9..036ed50b7 100644 --- a/src/torchcodec/_core/_ffmpeg_ops.py +++ b/src/torchcodec/_core/_ffmpeg_ops.py @@ -121,6 +121,9 @@ def add_video_stream( _blocks_packet_decoder_receive_frame = ( torch.ops.torchcodec_ns._blocks_packet_decoder_receive_frame.default ) +_blocks_audio_packet_decoder_receive_frame = ( + torch.ops.torchcodec_ns._blocks_audio_packet_decoder_receive_frame.default +) _blocks_create_color_converter = ( torch.ops.torchcodec_ns._blocks_create_color_converter.default ) diff --git a/src/torchcodec/_core/custom_ops.cpp b/src/torchcodec/_core/custom_ops.cpp index 483f0cd9f..5c0ae4f49 100644 --- a/src/torchcodec/_core/custom_ops.cpp +++ b/src/torchcodec/_core/custom_ops.cpp @@ -91,6 +91,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(torchcodec_ns, m) { m.def("_blocks_packet_decoder_reset(Tensor(a!) decoder) -> ()"); m.def( "_blocks_packet_decoder_receive_frame(Tensor(a!) decoder) -> (Tensor, int, float, float, Device, Tensor)"); + m.def( + "_blocks_audio_packet_decoder_receive_frame(Tensor(a!) decoder) -> (Tensor, int, float, float, int, str)"); m.def( "_blocks_create_color_converter(str device=\"cpu\", str output_dtype=\"uint8\") -> Tensor"); m.def( @@ -978,6 +980,46 @@ OpsReceiveFrameOutput _blocks_packet_decoder_receive_frame( storage); } +// (samples, status, pts_seconds, duration_seconds, sample_rate, +// sample_format). `samples` is [num_channels, num_samples] in the frame's own +// sample type; there is no frame handle because, unlike a video frame, nothing +// downstream needs the AVFrame itself. +using OpsReceiveAudioFrameOutput = std:: + tuple; + +OpsReceiveAudioFrameOutput _blocks_audio_packet_decoder_receive_frame( + torch::stable::Tensor& decoder) { + PacketDecoder* decoder_ptr = unwrap_tensor_to_pointer(decoder); + STD_TORCH_CHECK( + decoder_ptr->media_type() == AVMEDIA_TYPE_AUDIO, + "This PacketDecoder decodes video, not audio."); + + UniqueAVFrame av_frame(av_frame_alloc()); + STD_TORCH_CHECK(av_frame != nullptr, "Failed to allocate AVFrame"); + int status = decoder_ptr->receive_frame(av_frame); + if (status != AVSUCCESS) { + return std::make_tuple( + torch::stable::empty({int64_t(0)}, kStableUInt8), + static_cast(status), + 0.0, + 0.0, + static_cast(0), + std::string("")); + } + + AVRational time_base = decoder_ptr->time_base(); + const char* sample_format_name = + av_get_sample_fmt_name(static_cast(av_frame->format)); + return std::make_tuple( + audio_samples(*av_frame), + static_cast(0), + pts_to_seconds(get_pts_or_dts(*av_frame), time_base), + pts_to_seconds(get_duration(*av_frame), time_base), + static_cast(av_frame->sample_rate), + std::string( + sample_format_name == nullptr ? "unknown" : sample_format_name)); +} + torch::stable::Tensor _blocks_create_color_converter( std::string device, std::string output_dtype) { @@ -1613,6 +1655,9 @@ STABLE_TORCH_LIBRARY_IMPL(torchcodec_ns, CPU, m) { m.impl( "_blocks_packet_decoder_receive_frame", TORCH_BOX(&_blocks_packet_decoder_receive_frame)); + m.impl( + "_blocks_audio_packet_decoder_receive_frame", + TORCH_BOX(&_blocks_audio_packet_decoder_receive_frame)); m.impl("_blocks_convert_frame", TORCH_BOX(&_blocks_convert_frame)); m.impl("_blocks_frame_metadata", TORCH_BOX(&_blocks_frame_metadata)); m.impl("_blocks_frame_planes", TORCH_BOX(&_blocks_frame_planes)); diff --git a/src/torchcodec/_core/sources.bzl b/src/torchcodec/_core/sources.bzl index 6e2facadd..f1c9cb903 100644 --- a/src/torchcodec/_core/sources.bzl +++ b/src/torchcodec/_core/sources.bzl @@ -33,6 +33,7 @@ decoder_core_sources = [ "CpuDeviceInterface.cpp", "Demuxer.cpp", "PacketDecoder.cpp", + "AudioCommon.cpp", "ColorConverter.cpp", "SingleStreamDecoder.cpp", "Encoder.cpp", diff --git a/src/torchcodec/decoders/_blocks/__init__.py b/src/torchcodec/decoders/_blocks/__init__.py index 34fd5c45f..3126a9dc7 100644 --- a/src/torchcodec/decoders/_blocks/__init__.py +++ b/src/torchcodec/decoders/_blocks/__init__.py @@ -7,13 +7,14 @@ """Private, experimental building-block decode API. Exposes the three decode stages -- :class:`VideoDemuxer`, -:class:`PacketDecoder`, :class:`ColorConverter` -- as passive, composable, +:class:`VideoPacketDecoder`, :class:`ColorConverter` -- as passive, composable, GIL-releasing units, so a caller can build its own (threaded) decode pipeline and tune how the stages overlap. The blocks do no threading themselves. -:class:`AudioDemuxer` is the audio counterpart of :class:`VideoDemuxer`; -:class:`PacketDecoder` is shared, since decoding is the same operation either -way. +Audio has the same three stages: :class:`AudioDemuxer` and +:class:`AudioPacketDecoder`. The two decoders are a single class in C++, since +decoding is the same operation either way; they are separate here because what +they hand out, and how they are configured, isn't. This is experimental and private; the API may change. See API_breakdown_claude_plan.md for the design and rationale. @@ -21,16 +22,18 @@ from ._color_converter import ColorConverter from ._demuxer import AudioDemuxer, StreamIndex, VideoDemuxer -from ._frame import Packet, RawFrame -from ._packet_decoder import PacketDecoder +from ._frame import Packet, RawAudioSamples, RawFrame +from ._packet_decoder import AudioPacketDecoder, VideoPacketDecoder __all__ = [ "VideoDemuxer", "AudioDemuxer", - "PacketDecoder", + "VideoPacketDecoder", + "AudioPacketDecoder", "ColorConverter", "Packet", "RawFrame", + "RawAudioSamples", "StreamIndex", ] diff --git a/src/torchcodec/decoders/_blocks/_demuxer.py b/src/torchcodec/decoders/_blocks/_demuxer.py index dff25cfcb..b40323ed7 100644 --- a/src/torchcodec/decoders/_blocks/_demuxer.py +++ b/src/torchcodec/decoders/_blocks/_demuxer.py @@ -182,7 +182,7 @@ def seek(self, seconds: float) -> None: """Move the demuxer to ``seconds``. A seek invalidates whatever the decoder is holding on to, so the - :class:`PacketDecoder` must be ``reset()`` afterwards. + packet decoder must be ``reset()`` afterwards. Where you land, and what comes out first, depends on the medium. For video, a decoder can only start on a :term:`keyframe`, so this lands on @@ -212,8 +212,8 @@ class VideoDemuxer(_BaseDemuxer): of the file, or from wherever :meth:`seek` left it. A :class:`VideoDemuxer` also carries the stream configuration used to build a - :class:`PacketDecoder`, so that is constructed from a demuxer and no extra - container is opened. + :class:`VideoPacketDecoder`, so that is constructed from a demuxer and no + extra container is opened. Args: source (str, ``Pathlib.path``, bytes, ``torch.Tensor`` or file-like object): The source of the video: @@ -246,7 +246,7 @@ def scan(self) -> StreamIndex: The index always covers the entire stream, wherever the demuxer currently is, and the demuxer is left back at the start - so a - :class:`PacketDecoder` built from it must be ``reset()``, as after a + :class:`VideoPacketDecoder` built from it must be ``reset()``, as after a ``seek()``. Nothing is cached: calling this twice scans twice. """ pts, duration, is_key_frame, time_base_num, time_base_den = ( @@ -270,8 +270,8 @@ class AudioDemuxer(_BaseDemuxer): of the file, or from wherever :meth:`seek` left it. An :class:`AudioDemuxer` also carries the stream configuration used to build - a :class:`PacketDecoder`, so that is constructed from a demuxer and no extra - container is opened. + an :class:`AudioPacketDecoder`, so that is constructed from a demuxer and + no extra container is opened. Unlike :class:`VideoDemuxer` there is no ``scan()``: a :class:`StreamIndex` describes keyframes and frame indices, and audio has neither. diff --git a/src/torchcodec/decoders/_blocks/_frame.py b/src/torchcodec/decoders/_blocks/_frame.py index 5842dde1e..29b3d49cb 100644 --- a/src/torchcodec/decoders/_blocks/_frame.py +++ b/src/torchcodec/decoders/_blocks/_frame.py @@ -6,6 +6,7 @@ from __future__ import annotations +from dataclasses import dataclass from typing import NamedTuple import torch @@ -28,7 +29,7 @@ class _Metadata(NamedTuple): class Packet: """Opaque, thread-movable handle to a demuxed (compressed) packet. - Produced by :class:`VideoDemuxer`, consumed by :class:`PacketDecoder`. It wraps a raw + Produced by :class:`VideoDemuxer`, consumed by :class:`VideoPacketDecoder`. It wraps a raw pointer, so it is only valid within the process that created it (it cannot cross a process boundary). """ @@ -43,7 +44,7 @@ class RawFrame: """A decoded (YUV) frame, as the decoder produced it: an opaque, thread-movable handle to the frame plus everything describing it. - Produced by :class:`PacketDecoder`, consumed by :class:`ColorConverter`. The + Produced by :class:`VideoPacketDecoder`, consumed by :class:`ColorConverter`. The handle wraps a raw pointer and is process-local. ``pts_seconds`` and ``duration_seconds`` are stamped by the decoder (which knows the stream time base) and carried here so the :class:`ColorConverter` need not be bound to @@ -162,3 +163,51 @@ def planes(self) -> tuple[torch.Tensor, ...]: # views. self._planes = tuple(p for p in planes if p.numel() > 0) return self._planes + + +# TODO_API_BREAKDOWN DESIGN P1: API design - the class name, and whether +# sample_format is worth carrying now that the layout it describes has been +# normalized away. +@dataclass +class RawAudioSamples: + """One decoded audio frame's samples, as the decoder produced them. + + Produced by :class:`VideoPacketDecoder` for an :class:`AudioDemuxer`'s stream, + consumed by :class:`AudioConverter`. This is the audio counterpart of + :class:`RawFrame`, and like it, nothing has been converted: the samples are + in the codec's own sample type. + + It is not a handle, unlike :class:`RawFrame`. Audio frames are a few kB, so + the samples are copied out of the ``AVFrame`` rather than viewed, which + also lets ``[num_channels, num_samples]`` be the layout for every format: + planar ones store each channel in its own allocation and packed ones + interleave them, so neither is that shape as it stands. + + Attributes: + data (torch.Tensor): ``[num_channels, num_samples]``, in the dtype that + holds the source's samples exactly: ``uint8`` for ``u8``, ``int16`` + for ``s16``, ``int32`` for ``s32``, ``float32`` for ``flt``, + ``float64`` for ``dbl``. Note the integer ones are *not* normalized + to ``[-1, 1]``; :class:`AudioConverter` is what does that. + sample_rate (int): The source's sample rate, in Hz. + sample_format (str): FFmpeg sample-format name, e.g. ``"s16p"`` or + ``"fltp"``. This is the format the samples were decoded in, kept + for provenance: the trailing ``p`` (planar) no longer describes + :attr:`data`, whose layout is always the same. + pts_seconds (float): Presentation timestamp of the first sample. + duration_seconds (float): How long these samples last. + """ + + data: torch.Tensor + sample_rate: int + sample_format: str + pts_seconds: float + duration_seconds: float + + @property + def num_channels(self) -> int: + return self.data.shape[0] + + @property + def num_samples(self) -> int: + return self.data.shape[1] diff --git a/src/torchcodec/decoders/_blocks/_packet_decoder.py b/src/torchcodec/decoders/_blocks/_packet_decoder.py index 82916c942..8b7457b86 100644 --- a/src/torchcodec/decoders/_blocks/_packet_decoder.py +++ b/src/torchcodec/decoders/_blocks/_packet_decoder.py @@ -6,9 +6,12 @@ from __future__ import annotations +from typing import Generic, TypeVar + import torch from torchcodec._core.ops import ( + _blocks_audio_packet_decoder_receive_frame, _blocks_create_packet_decoder, _blocks_packet_decoder_receive_frame, _blocks_packet_decoder_reset, @@ -17,32 +20,78 @@ ) from .._decoder_utils import convert_device_to_str -from ._demuxer import VideoDemuxer -from ._frame import Packet, RawFrame +from ._demuxer import AudioDemuxer, VideoDemuxer +from ._frame import Packet, RawAudioSamples, RawFrame # TODO_API_BREAKDOWN DOC P1 revisit every single docstring / comments at some point. +_Decoded = TypeVar("_Decoded", RawFrame, RawAudioSamples) + + +class _BasePacketDecoder(Generic[_Decoded]): + """Shared machinery for :class:`VideoPacketDecoder` and + :class:`AudioPacketDecoder`. + + Decoding is the same ``avcodec_send_packet`` / ``avcodec_receive_frame`` + pair for both, so the C++ side is a single class; what differs is only what + a decoded frame is turned into, which is :meth:`_receive_ready_frames`. + """ + + def __init__(self, demuxer, device_str: str): + self._handle = _blocks_create_packet_decoder( + demuxer._handle, num_threads=1, device=device_str + ) + self._drained = False + + def _receive_ready_frames(self) -> list[_Decoded]: + raise NotImplementedError + + def decode(self, packet: Packet) -> list[_Decoded]: + """Send one packet and return whatever is now ready (possibly empty, + e.g. while the codec buffers B-frames).""" + if self._drained: + raise RuntimeError( + "This decoder has been drained, and a codec that has been told " + "the stream ended ignores any further packet. Create a new " + "decoder to decode another stream." + ) + status = _blocks_packet_decoder_send_packet(self._handle, packet._handle) + if status < 0: + raise RuntimeError(f"Failed to send packet to decoder (status {status})") + return self._receive_ready_frames() + + def drain(self) -> list[_Decoded]: + """Tell the codec the stream ended, and return the frames it was still + holding on to.""" + _blocks_packet_decoder_send_eof(self._handle) + frames = self._receive_ready_frames() + self._drained = True + return frames + + def reset(self) -> None: + """Drop the codec's buffered state and start over. Needed after the + demuxer seeked, and after ``drain()``.""" + _blocks_packet_decoder_reset(self._handle) + self._drained = False + -class PacketDecoder: +class VideoPacketDecoder(_BasePacketDecoder[RawFrame]): """Decode building block: turns compressed :class:`Packet`\\ s into decoded (YUV) :class:`RawFrame`\\ s. - Built from a :class:`VideoDemuxer` (for its codec parameters) and stateful (it - holds the codec's reference-frame buffer). Passive and *not* thread-safe: - use one ``PacketDecoder`` per thread. FFmpeg's internal codec thread count - is kept at 1 for now (not exposed); parallelism comes from composing blocks - on your own threads. + Built from a :class:`VideoDemuxer` (for its codec parameters) and stateful + (it holds the codec's reference-frame buffer). Passive and *not* + thread-safe: use one ``VideoPacketDecoder`` per thread. FFmpeg's internal + codec thread count is kept at 1 for now (not exposed); parallelism comes + from composing blocks on your own threads. ``device`` accepts a string or a ``torch.device``. It defaults to ``None``, which means the current default device (see ``torch.set_default_device``). """ def __init__(self, demuxer: VideoDemuxer, device: str | torch.device | None = None): - self._handle = _blocks_create_packet_decoder( - demuxer._handle, num_threads=1, device=convert_device_to_str(device) - ) - self._drained = False + super().__init__(demuxer, convert_device_to_str(device)) def _receive_ready_frames(self) -> list[RawFrame]: frames = [] @@ -63,26 +112,39 @@ def _receive_ready_frames(self) -> list[RawFrame]: ) return frames - def decode(self, packet: Packet) -> list[RawFrame]: - """Send one packet and return whatever frames are now ready (possibly - empty, e.g. while the codec buffers B-frames).""" - if self._drained: - raise RuntimeError( - "This PacketDecoder has been drained, and a codec that has been " - "told the stream ended ignores any further packet. Create a new " - "PacketDecoder to decode another stream." - ) - status = _blocks_packet_decoder_send_packet(self._handle, packet._handle) - if status < 0: - raise RuntimeError(f"Failed to send packet to decoder (status {status})") - return self._receive_ready_frames() - def drain(self) -> list[RawFrame]: - _blocks_packet_decoder_send_eof(self._handle) - frames = self._receive_ready_frames() - self._drained = True - return frames +class AudioPacketDecoder(_BasePacketDecoder[RawAudioSamples]): + """Decode building block: turns compressed :class:`Packet`\\ s into + :class:`RawAudioSamples`. - def reset(self) -> None: - _blocks_packet_decoder_reset(self._handle) - self._drained = False + Built from an :class:`AudioDemuxer` (for its codec parameters) and stateful: + a lossy codec's overlap-add state means the frames decoded right after a + seek are subtly wrong until it re-primes, so ``reset()`` is necessary but + not by itself sufficient - see :meth:`AudioDemuxer.seek`. Passive and *not* + thread-safe: use one ``AudioPacketDecoder`` per thread. + + There is no ``device`` parameter: audio is always decoded on the CPU, and + that doesn't change with ``torch.set_default_device``. + """ + + def __init__(self, demuxer: AudioDemuxer): + super().__init__(demuxer, "cpu") + + def _receive_ready_frames(self) -> list[RawAudioSamples]: + samples = [] + while True: + data, status, pts_seconds, duration_seconds, sample_rate, sample_format = ( + _blocks_audio_packet_decoder_receive_frame(self._handle) + ) + if status != 0: # EAGAIN (need more packets) or EOF: nothing ready + break + samples.append( + RawAudioSamples( + data=data, + sample_rate=sample_rate, + sample_format=sample_format, + pts_seconds=pts_seconds, + duration_seconds=duration_seconds, + ) + ) + return samples diff --git a/test/test_decoders.py b/test/test_decoders.py index d4605343b..914e3c1ce 100644 --- a/test/test_decoders.py +++ b/test/test_decoders.py @@ -43,11 +43,13 @@ ) from torchcodec.decoders._blocks import ( AudioDemuxer, + AudioPacketDecoder, ColorConverter, Packet, - PacketDecoder, + RawAudioSamples, RawFrame, VideoDemuxer, + VideoPacketDecoder, ) from torchcodec.decoders._decoder_utils import _get_cuda_backend from torchcodec.decoders._image_decoders import _source_to_tensor @@ -3635,7 +3637,7 @@ class TestBlocks: @pytest.mark.parametrize("device", _block_devices()) def test_block_output_types(self, device): - # VideoDemuxer yields Packets, PacketDecoder yields RawFrames, and + # VideoDemuxer yields Packets, VideoPacketDecoder yields RawFrames, and # ColorConverter yields Frames with the expected shape/dtype. demuxer, decoder, converter = self._make_blocks(NASA_VIDEO.path, device) @@ -3710,7 +3712,7 @@ def drain(): @staticmethod def _make_blocks(path, device): demuxer = VideoDemuxer(path) - decoder = PacketDecoder(demuxer, device=device) + decoder = VideoPacketDecoder(demuxer, device=device) converter = ColorConverter(device=device) return demuxer, decoder, converter @@ -3761,7 +3763,7 @@ def _to_frame_batch(self, frames): def _assert_matches_video_decoder(got, ref, video): # We typically want exact equality but cannot achieve it on CUDA for HDR # videos that are downscaled to uint8: the VideoDecoder will ask NVDEC - # to output an 8bit surface while the PacketDecoder will decode on the + # to output an 8bit surface while the VideoPacketDecoder will decode on the # 16bit surface (by contract), so there are minor differences. if got.is_cuda and got.dtype == torch.uint8 and video in _HDR_VIDEOS: assert_tensor_close_on_at_least(got, ref, percentage=99, atol=2) @@ -3967,14 +3969,14 @@ def _first_frame(self, path, device): @pytest.mark.parametrize("device_str", _block_devices()) def test_device_none_default_device(self, device_str): - # PacketDecoder and ColorConverter default to device=None, which should + # VideoPacketDecoder and ColorConverter default to device=None, which should # respect both the torch.device() context manager and # torch.set_default_device(). def assert_first_frame_is_on_default_device(): # Note the absence of any device parameter. demuxer = VideoDemuxer(NASA_VIDEO.path) - decoder = PacketDecoder(demuxer) + decoder = VideoPacketDecoder(demuxer) converter = ColorConverter() decoded = next(self._decode(decoder, self._demux(demuxer))) assert decoded.planes[0].device.type == device_str @@ -4078,7 +4080,7 @@ def test_no_rotation(self, device): @pytest.mark.needs_cuda @pytest.mark.parametrize("record_stream", (True, False)) def test_storage_record_stream(self, record_stream): - # Using PacketDecoder on one stream and consuming the frames on a + # Using VideoPacketDecoder on one stream and consuming the frames on a # different stream requires the user to call record_stream() on the # frame storage. # Without the record_stream() call the decoder's next frame may be @@ -4281,7 +4283,7 @@ def test_neutral_chroma_is_grayscale(self, video, device): ), ) def test_cpu_fallback_is_on_cuda(self, video, expected_pix_fmt): - # A CUDA PacketDecoder hands out CUDA frames even for the streams it has + # A CUDA VideoPacketDecoder hands out CUDA frames even for the streams it has # to decode on the CPU, and they're in an NVDEC surface format like any # other CUDA frame. assert VideoDecoder(video.path, device="cuda").cpu_fallback @@ -4478,7 +4480,7 @@ def test_seek_without_reset_yields_stale_frames(self, device): seconds = video_decoder.get_frame_at(keyframe_index).pts_seconds demuxer = VideoDemuxer(NASA_VIDEO.path) - decoder = PacketDecoder(demuxer, device=device) + decoder = VideoPacketDecoder(demuxer, device=device) num_decoded = 0 for packet in demuxer: # decode a bit, so frames pile up in the codec num_decoded += len(decoder.decode(packet)) @@ -4743,7 +4745,7 @@ def test_seek_on_non_seekable_source_raises(self, tmp_path): ) try: demuxer = VideoDemuxer(fifo_path) - decoder = PacketDecoder(demuxer) + decoder = VideoPacketDecoder(demuxer) num_decoded = 0 for packet in demuxer: # make sure the stream is really flowing num_decoded += len(decoder.decode(packet)) @@ -4764,7 +4766,7 @@ def test_decode_after_drain_raises(self, device): # ignores anything sent afterwards. Rather than silently decoding # nothing, say so. demuxer = VideoDemuxer(H265_VIDEO.path) - decoder = PacketDecoder(demuxer, device=device) + decoder = VideoPacketDecoder(demuxer, device=device) packet = demuxer.next_packet() decoder.decode(packet) decoder.drain() @@ -4926,7 +4928,7 @@ def test_scan_leaves_demuxer_at_the_start(self, video, device): demuxer = VideoDemuxer(video.path) index = demuxer.scan() - decoder = PacketDecoder(demuxer, device=device) + decoder = VideoPacketDecoder(demuxer, device=device) converter = ColorConverter(device=device) frames = list( self._convert(converter, self._decode(decoder, self._demux(demuxer))) @@ -4980,7 +4982,7 @@ def test_stream_index(self, stream_index): # nasa_13013.mp4 has two video streams, 0 and 3, of different sizes, # and 3 is the best one, i.e. the one used when nothing is requested. demuxer = VideoDemuxer(NASA_VIDEO.path, stream_index=stream_index) - decoder = PacketDecoder(demuxer) + decoder = VideoPacketDecoder(demuxer) converter = ColorConverter() got = [ converter.convert(raw_frame) @@ -5057,19 +5059,185 @@ def test_audio_demuxer_seek(self): assert 0 < num_packets_after_seek < num_packets_from_start + # ===== Audio decoding: RawAudioSamples ===== + + @staticmethod + def _decode_audio(asset, stream_index=None, seek_seconds=None): + demuxer = AudioDemuxer(asset.path, stream_index=stream_index) + decoder = AudioPacketDecoder(demuxer) + if seek_seconds is not None: + demuxer.seek(seek_seconds) + decoder.reset() + chunks = [] + for packet in demuxer: + chunks += decoder.decode(packet) + chunks += decoder.drain() + return chunks + + @pytest.mark.parametrize( + "asset, sample_format, dtype", + ( + (SINE_MONO_U8, "u8", torch.uint8), + (SINE_MONO_S16, "s16", torch.int16), + (SINE_MONO_S32, "s32", torch.int32), + # FFmpeg has no 24-bit sample format, so a 24-bit source is s32. + (SINE_MONO_S24, "s32", torch.int32), + (SINE_MONO_F32, "flt", torch.float32), + (SINE_MONO_F64, "dbl", torch.float64), + (SINE_STEREO_MP2_MPEG_PS, "s16p", torch.int16), + (SINE_16_CHANNEL_S16, "s16", torch.int16), + (NASA_AUDIO_MP3, "fltp", torch.float32), + ), + ) + def test_audio_raw_samples_dtype_and_shape(self, asset, sample_format, dtype): + # The decoder hands out the codec's own sample type, always as + # [num_channels, num_samples]. The packed formats above (no trailing + # 'p') are the ones exercising the de-interleaving path. + chunks = self._decode_audio(asset) + assert len(chunks) > 0 + + for chunk in chunks: + assert isinstance(chunk, RawAudioSamples) + assert chunk.sample_format == sample_format + assert chunk.data.dtype == dtype + assert chunk.data.ndim == 2 + assert chunk.data.is_contiguous() + assert chunk.num_channels == asset.num_channels + assert chunk.sample_rate == asset.sample_rate + assert chunk.duration_seconds >= 0 + + @pytest.mark.parametrize( + "asset", + ( + SINE_MONO_U8, + SINE_MONO_S16, + SINE_MONO_S32, + SINE_MONO_F32, + SINE_MONO_F64, + SINE_STEREO_MP2_MPEG_PS, + SINE_16_CHANNEL_S16, + NASA_AUDIO_MP3, + NASA_AUDIO, + ), + ) + @pytest.mark.parametrize("seek_fraction", (None, 1 / 3, 2 / 3)) + def test_audio_raw_samples_match_audio_decoder(self, asset, seek_fraction): + # We hand out the true source samples: normalizing them the way FFmpeg + # does reproduces AudioDecoder's output bit for bit. This is also what + # pins the de-interleaving, most visibly on the 16-channel asset. + # + # It holds after a seek too, but only once the caller has done the + # pre-roll these blocks don't do: a lossy codec decodes its first frames + # after a seek from a flushed state, so they come out subtly wrong - + # plausible, but not what whole-file decoding gives - until it + # re-primes. Dropping those frames is exactly what pre-rolling means, + # and everything from there on is bit exact again. Without the drop, + # mp3 and aac diverge over their first ~1000-1600 samples and match + # perfectly after that. + if seek_fraction is not None and asset is SINE_STEREO_MP2_MPEG_PS: + pytest.skip( + "MPEG-PS resync after a seek is unreliable for both the blocks " + "and AudioDecoder (which raises seeking this file to 2.8s), so " + "it can't tell us anything here. See " + "test_audio_decoder_mpeg_ps_resync_after_seek." + ) + + seek_seconds = ( + None if seek_fraction is None else asset.duration_seconds * seek_fraction + ) + chunks = self._decode_audio(asset, seek_seconds=seek_seconds) + if seek_seconds is not None: + # Same number of frames SingleStreamDecoder pre-rolls by, see + # Note [Audio pre-roll and post-roll]. + chunks = chunks[4:] + assert len(chunks) > 0 + + raw = torch.cat([chunk.data for chunk in chunks], dim=1) + + if raw.dtype == torch.uint8: + got = (raw.to(torch.float32) - 128) / 128 + elif raw.dtype in (torch.int16, torch.int32): + got = raw.to(torch.float32) / -float(torch.iinfo(raw.dtype).min) + else: + got = raw.to(torch.float32) + + decoder = AudioDecoder(asset.path) + if seek_seconds is None: + expected = decoder.get_all_samples().data + else: + # Re-anchor on the first frame we kept, so both sides start on the + # same sample. + expected = decoder.get_samples_played_in_range( + start_seconds=chunks[0].pts_seconds + ).data + torch.testing.assert_close(got, expected, atol=0, rtol=0) + + def test_audio_raw_samples_pts(self): + chunks = self._decode_audio(SINE_MONO_S16) + pts = [chunk.pts_seconds for chunk in chunks] + assert pts == sorted(pts) + assert pts[0] == pytest.approx(0, abs=1e-6) + + def test_decoder_output_type_follows_the_demuxer(self): + # The two decoders are one class in C++; the split is a Python-level + # one, so that each has an exact output type and its own arguments. + for demuxer_class, decoder_class, expected_type in ( + (AudioDemuxer, AudioPacketDecoder, RawAudioSamples), + (VideoDemuxer, VideoPacketDecoder, RawFrame), + ): + demuxer = demuxer_class(NASA_VIDEO.path) + decoder = decoder_class(demuxer) + # A codec needs more than one packet before it outputs anything. + decoded = [] + while not decoded: + decoded = decoder.decode(demuxer.next_packet()) + assert isinstance(decoded[0], expected_type) + + def test_audio_decoder_takes_no_device(self): + with pytest.raises(TypeError, match="device"): + AudioPacketDecoder(AudioDemuxer(NASA_AUDIO_MP3.path), device="cuda") + + def test_audio_decoder_mpeg_ps_resync_after_seek(self): + # Seeking an MPEG program stream lands on a container-level byte + # offset, so the parser resumes mid-frame and the packets it rebuilds + # are AVERROR_INVALIDDATA until it resyncs. That's a property of the + # container, not of the codec: it applies to this file's audio stream + # exactly as it does to a video one. Without the resync handling this + # raises "Failed to send packet to decoder" on the very first packet. + asset = SINE_STEREO_MP2_MPEG_PS + demuxer = AudioDemuxer(asset.path) + decoder = AudioPacketDecoder(demuxer) + demuxer.seek(asset.duration_seconds / 2) + decoder.reset() + + chunks = [] + for packet in demuxer: + chunks += decoder.decode(packet) + chunks += decoder.drain() + + assert len(chunks) > 0 + num_samples = sum(chunk.num_samples for chunk in chunks) + assert 0 < num_samples < asset.duration_seconds * asset.sample_rate + + @pytest.mark.parametrize( + "demuxer_class", (VideoDemuxer, AudioDemuxer), ids=("video", "audio") + ) @pytest.mark.parametrize("stream_index", (-1, 6, 1000)) - def test_invalid_stream_index_raises(self, stream_index): + def test_invalid_stream_index_raises(self, demuxer_class, stream_index): with pytest.raises(RuntimeError, match="is not a valid stream"): - VideoDemuxer(NASA_VIDEO.path, stream_index=stream_index) + demuxer_class(NASA_VIDEO.path, stream_index=stream_index) - def test_bad_source_type_raises(self): + @pytest.mark.parametrize( + "demuxer_class", (VideoDemuxer, AudioDemuxer), ids=("video", "audio") + ) + def test_bad_source_type_raises(self, demuxer_class): with pytest.raises(TypeError, match="Unknown source type"): - VideoDemuxer(123) + demuxer_class(123) # user mistakenly forgets to specify binary reading when creating a # file-like object from open() with pytest.raises(TypeError, match="binary reading?"): - VideoDemuxer(open(NASA_VIDEO.path)) + demuxer_class(open(NASA_VIDEO.path)) # Small helpers to avoid having to always specify the same skip marks and decode_fn