From b736d389e71ae10a951e5607b27b74c73477c1b9 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Thu, 21 May 2026 16:08:08 +0300 Subject: [PATCH] init --- src/torchcodec/decoders/_video_decoder.py | 59 ++++++++++++++++++++++- test/test_decoders.py | 59 +++++++++++++++++++++++ 2 files changed, 116 insertions(+), 2 deletions(-) diff --git a/src/torchcodec/decoders/_video_decoder.py b/src/torchcodec/decoders/_video_decoder.py index b7e7c7f98..7fa076bc5 100644 --- a/src/torchcodec/decoders/_video_decoder.py +++ b/src/torchcodec/decoders/_video_decoder.py @@ -3,7 +3,7 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. - +from __future__ import annotations import io import json @@ -11,7 +11,7 @@ from collections.abc import Sequence from dataclasses import dataclass, field from pathlib import Path -from typing import Literal +from typing import Literal, TYPE_CHECKING import torch from torch import device as torch_device, nn, Tensor @@ -21,6 +21,9 @@ from torchcodec.decoders._decoder_utils import _get_cuda_backend from torchcodec.transforms import DecoderTransform +if TYPE_CHECKING: + from torchcodec.decoders._audio_decoder import AudioDecoder + @dataclass class CpuFallbackStatus: @@ -146,6 +149,13 @@ class VideoDecoder: Alternative field names "pkt_pts" and "pkt_duration" are also supported. Read more about this parameter in: :ref:`sphx_glr_generated_examples_decoding_custom_frame_mappings.py` + audio_stream_index (int, optional): Stream index of the audio stream to + use when :attr:`audio` is accessed. If unspecified, the + :term:`best stream` is used. Ignored for ``torch.Tensor`` sources. + audio_sample_rate (int, optional): Desired output sample rate for the + :attr:`audio` decoder. Defaults to the stream's native sample rate. + audio_num_channels (int, optional): Desired number of channels for the + :attr:`audio` decoder. Defaults to the stream's native channel count. Attributes: metadata (VideoStreamMetadata): Metadata of the video stream. @@ -171,8 +181,22 @@ def __init__( custom_frame_mappings: ( str | bytes | io.RawIOBase | io.BufferedReader | None ) = None, + audio_stream_index: int | None = None, + audio_sample_rate: int | None = None, + audio_num_channels: int | None = None, ): torch._C._log_api_usage_once("torchcodec.decoders.VideoDecoder") + + # Normalize file-like objects to bytes so the source can be stored and + # later passed to AudioDecoder independently (file-likes have a single + # shared cursor that two decoders would stomp on each other). + if not isinstance(source, (str, Path, bytes, Tensor)) and hasattr( + source, "read" + ): + if hasattr(source, "seek"): + source.seek(0) + source = source.read() + allowed_seek_modes = ("exact", "approximate") if seek_mode not in allowed_seek_modes: raise ValueError( @@ -245,6 +269,13 @@ def __init__( else: self._cpu_fallback._backend = "CPU" + self._source = source + self._audio_stream_index = audio_stream_index + self._audio_sample_rate = audio_sample_rate + self._audio_num_channels = audio_num_channels + self._audio: AudioDecoder | None = None + self._audio_initialized = False + def __len__(self) -> int: return self._num_frames @@ -275,6 +306,30 @@ def cpu_fallback(self) -> CpuFallbackStatus: return self._cpu_fallback + @property + def audio(self) -> AudioDecoder | None: + """An :class:`AudioDecoder` for the best audio stream in the same + source, or ``None`` if the source has no audio stream or is a + ``torch.Tensor`` (raw encoded bytes with no container). + + The decoder is created lazily on first access and cached. + """ + if not self._audio_initialized: + if not isinstance(self._source, Tensor): + try: + from torchcodec.decoders._audio_decoder import AudioDecoder + + self._audio = AudioDecoder( + self._source, + stream_index=self._audio_stream_index, + sample_rate=self._audio_sample_rate, + num_channels=self._audio_num_channels, + ) + except ValueError: + pass + self._audio_initialized = True + return self._audio + def _getitem_int(self, key: int) -> Tensor: assert isinstance(key, int) diff --git a/test/test_decoders.py b/test/test_decoders.py index bb2d65e1f..579450aed 100644 --- a/test/test_decoders.py +++ b/test/test_decoders.py @@ -2816,3 +2816,62 @@ def test_get_samples_played_in_range_errors(self): match="No samples to decode. This is probably because start_seconds is too high\\(10\\)", ): wav_dec.get_samples_played_in_range(10.0, 12.0) + + +class TestVideoDecoderAudio: + def test_audio_property_returns_audio_decoder(self): + decoder = VideoDecoder(NASA_VIDEO.path) + assert isinstance(decoder.audio, AudioDecoder) + + def test_audio_metadata(self): + decoder = VideoDecoder(NASA_VIDEO.path) + assert decoder.audio is not None + assert isinstance(decoder.audio.metadata, AudioStreamMetadata) + assert decoder.audio.stream_index == NASA_AUDIO.default_stream_index + + def test_audio_samples_access(self): + decoder = VideoDecoder(NASA_VIDEO.path) + assert decoder.audio is not None + samples = decoder.audio.get_all_samples() + assert samples.data.ndim == 2 # (num_channels, num_samples) + assert samples.data.shape[0] == NASA_AUDIO.num_channels + + def test_audio_is_none_for_video_only_source(self): + decoder = VideoDecoder(H265_VIDEO.path) + assert decoder.audio is None + + def test_audio_is_cached(self): + decoder = VideoDecoder(NASA_VIDEO.path) + assert decoder.audio is decoder.audio + + def test_audio_is_none_for_tensor_source(self): + data = torch.frombuffer(NASA_VIDEO.path.read_bytes(), dtype=torch.uint8) + decoder = VideoDecoder(data) + assert decoder.audio is None + + def test_audio_for_bytes_source(self): + decoder = VideoDecoder(NASA_VIDEO.path.read_bytes()) + assert isinstance(decoder.audio, AudioDecoder) + + def test_audio_for_file_like_source(self): + with open(NASA_VIDEO.path, "rb") as f: + decoder = VideoDecoder(f) + assert isinstance(decoder.audio, AudioDecoder) + + def test_audio_stream_index_kwarg(self): + decoder = VideoDecoder(NASA_VIDEO.path, audio_stream_index=NASA_AUDIO.default_stream_index) + assert decoder.audio is not None + assert decoder.audio.stream_index == NASA_AUDIO.default_stream_index + + def test_audio_sample_rate_kwarg(self): + target_rate = 8_000 + decoder = VideoDecoder(NASA_VIDEO.path, audio_sample_rate=target_rate) + assert decoder.audio is not None + samples = decoder.audio.get_all_samples() + assert samples.sample_rate == target_rate + + def test_audio_num_channels_kwarg(self): + decoder = VideoDecoder(NASA_VIDEO.path, audio_num_channels=1) + assert decoder.audio is not None + samples = decoder.audio.get_all_samples() + assert samples.data.shape[0] == 1