# 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.
import io
from pathlib import Path
import torch
from torch import Tensor
from torchcodec import _core as core, AudioSamples
from torchcodec._core._decoder_utils import create_audio_decoder
[docs]
class AudioDecoder:
"""A single-stream audio decoder.
This can be used to decode audio from pure audio files (e.g. mp3, wav,
etc.), or from videos that contain audio streams (e.g. mp4 videos).
.. tip::
If you are decoding WAV files, consider using
:class:`~torchcodec.decoders.WavDecoder` instead for significantly faster
decoding.
Returned samples are float samples normalized in [-1, 1]
Args:
source (str, ``Pathlib.path``, bytes, ``torch.Tensor`` or file-like
object): The source of the video or audio:
- If ``str``: a local path or a URL to a video or audio file.
- If ``Pathlib.path``: a path to a local video or audio file.
- If ``bytes`` object or ``torch.Tensor``: the raw encoded audio data.
- If file-like object: we read video data from the object on demand. The object must
expose the methods `read(self, size: int) -> bytes` and
`seek(self, offset: int, whence: int) -> int`. Read more in:
:ref:`sphx_glr_generated_examples_decoding_file_like.py`.
stream_index (int, optional): Specifies which stream in the file to decode samples from.
Note that this index is absolute across all media types. If left unspecified, then
the :term:`best stream` is used.
sample_rate (int, optional): The desired output sample rate of the decoded samples.
By default, the sample rate of the source is used.
num_channels (int, optional): The desired number of channels of the decoded samples.
By default, the number of channels of the source is used.
Attributes:
metadata (AudioStreamMetadata): Metadata of the audio stream.
stream_index (int): The stream index that this decoder is retrieving samples from. If a
stream index was provided at initialization, this is the same value. If it was left
unspecified, this is the :term:`best stream`.
"""
def __init__(
self,
source: str | Path | io.RawIOBase | io.BufferedReader | bytes | Tensor,
*,
stream_index: int | None = None,
sample_rate: int | None = None,
num_channels: int | None = None,
):
torch._C._log_api_usage_once("torchcodec.decoders.AudioDecoder")
(
self._decoder,
self.stream_index,
self.metadata,
) = create_audio_decoder(
source=source,
seek_mode="approximate",
stream_index=stream_index,
sample_rate=sample_rate,
num_channels=num_channels,
)
self._desired_sample_rate = (
sample_rate if sample_rate is not None else self.metadata.sample_rate
)
[docs]
def get_all_samples(self) -> AudioSamples:
"""Returns all the audio samples from the source.
To decode samples in a specific range, use
:meth:`~torchcodec.decoders.AudioDecoder.get_samples_played_in_range`.
Returns:
AudioSamples: The samples within the file.
"""
return self.get_samples_played_in_range()
[docs]
def get_samples_played_in_range(
self, start_seconds: float = 0.0, stop_seconds: float | None = None
) -> AudioSamples:
"""Returns audio samples in the given range.
Samples are in the half open range [start_seconds, stop_seconds).
To decode all the samples from beginning to end, you can call this
method while leaving ``start_seconds`` and ``stop_seconds`` to their
default values, or use
:meth:`~torchcodec.decoders.AudioDecoder.get_all_samples` as a more
convenient alias.
Args:
start_seconds (float): Time, in seconds, of the start of the
range. Default: 0.
stop_seconds (float or None): Time, in seconds, of the end of the
range. As a half open range, the end is excluded. Default: None,
which decodes samples until the end.
Returns:
AudioSamples: The samples within the specified range.
"""
if stop_seconds is not None and not start_seconds <= stop_seconds:
raise ValueError(
f"Invalid start seconds: {start_seconds}. It must be less than or equal to stop seconds ({stop_seconds})."
)
frames, first_pts = core.get_frames_by_pts_in_range_audio(
self._decoder,
start_seconds=start_seconds,
stop_seconds=stop_seconds,
)
first_pts = first_pts.item()
# x = frame boundaries
#
# first_pts last_pts
# v v
# ....x..........x..........x...........x..........x..........x.....
# ^ ^
# start_seconds stop_seconds
#
# We want to return the samples in [start_seconds, stop_seconds). But
# because the core API is based on frames, the `frames` tensor contains
# the samples in [first_pts, last_pts)
# So we do some basic math to figure out the position of the view that
# we'll return.
sample_rate = self._desired_sample_rate
# TODO: metadata's sample_rate should probably not be Optional
assert sample_rate is not None # mypy.
if first_pts < start_seconds:
offset_beginning = round((start_seconds - first_pts) * sample_rate)
output_pts_seconds = start_seconds
else:
# In normal cases we'll have first_pts <= start_pts, but in some
# edge cases it's possible to have first_pts > start_seconds,
# typically if the stream's first frame's pts isn't exactly 0.
offset_beginning = 0
output_pts_seconds = first_pts
num_samples = frames.shape[1]
last_pts = first_pts + num_samples / sample_rate
if stop_seconds is not None and stop_seconds < last_pts:
offset_end = num_samples - round((last_pts - stop_seconds) * sample_rate)
else:
offset_end = num_samples
data = frames[:, offset_beginning:offset_end]
return AudioSamples(
data=data,
pts_seconds=output_pts_seconds,
duration_seconds=data.shape[1] / sample_rate,
sample_rate=sample_rate,
)