Add audio property to VideoDecoder

src/torchcodec/decoders/_video_decoder.py

@@ -3,15 +3,15 @@
 #
 # 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
 import numbers
 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)
 

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