torchcodec 0.10.0__cp312-cp312-manylinux_2_28_x86_64.whl

torchcodec/decoders/_audio_decoder.py

@@ -0,0 +1,185 @@
+# 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.decoders._decoder_utils import (
+    create_decoder,
+    ERROR_REPORTING_INSTRUCTIONS,
+)
+
+
+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).
+
+    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 = create_decoder(source=source, seek_mode="approximate")
+
+        container_metadata = core.get_container_metadata(self._decoder)
+        self.stream_index = (
+            container_metadata.best_audio_stream_index
+            if stream_index is None
+            else stream_index
+        )
+        if self.stream_index is None:
+            raise ValueError(
+                "The best audio stream is unknown and there is no specified stream. "
+                + ERROR_REPORTING_INSTRUCTIONS
+            )
+        if self.stream_index >= len(container_metadata.streams):
+            raise ValueError(
+                f"The stream at index {stream_index} is not a valid stream."
+            )
+
+        self.metadata = container_metadata.streams[self.stream_index]
+        if not isinstance(self.metadata, core._metadata.AudioStreamMetadata):
+            raise ValueError(
+                f"The stream at index {stream_index} is not an audio stream. "
+            )
+
+        self._desired_sample_rate = (
+            sample_rate if sample_rate is not None else self.metadata.sample_rate
+        )
+
+        core.add_audio_stream(
+            self._decoder,
+            stream_index=stream_index,
+            sample_rate=sample_rate,
+            num_channels=num_channels,
+        )
+
+    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()
+
+    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,
+        )

torchcodec/decoders/_decoder_utils.py

@@ -0,0 +1,113 @@
+# 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 contextvars
+import io
+
+from collections.abc import Generator
+from contextlib import contextmanager
+from pathlib import Path
+
+from torch import Tensor
+from torchcodec import _core as core
+
+ERROR_REPORTING_INSTRUCTIONS = """
+This should never happen. Please report an issue following the steps in
+https://github.com/pytorch/torchcodec/issues/new?assignees=&labels=&projects=&template=bug-report.yml.
+"""
+
+
+def create_decoder(
+    *,
+    source: str | Path | io.RawIOBase | io.BufferedReader | bytes | Tensor,
+    seek_mode: str,
+) -> Tensor:
+    if isinstance(source, str):
+        return core.create_from_file(source, seek_mode)
+    elif isinstance(source, Path):
+        return core.create_from_file(str(source), seek_mode)
+    elif isinstance(source, io.RawIOBase) or isinstance(source, io.BufferedReader):
+        return core.create_from_file_like(source, seek_mode)
+    elif isinstance(source, bytes):
+        return core.create_from_bytes(source, seek_mode)
+    elif isinstance(source, Tensor):
+        return core.create_from_tensor(source, seek_mode)
+    elif isinstance(source, io.TextIOBase):
+        raise TypeError(
+            "source is for reading text, likely from open(..., 'r'). Try with 'rb' for binary reading?"
+        )
+    elif hasattr(source, "read") and hasattr(source, "seek"):
+        # This check must be after checking for text-based reading. Also placing
+        # it last in general to be defensive: hasattr is a blunt instrument. We
+        # could use the inspect module to check for methods with the right
+        # signature.
+        return core.create_from_file_like(source, seek_mode)
+
+    raise TypeError(
+        f"Unknown source type: {type(source)}. "
+        "Supported types are str, Path, bytes, Tensor and file-like objects with "
+        "read(self, size: int) -> bytes and "
+        "seek(self, offset: int, whence: int) -> int methods."
+    )
+
+
+# Thread-local and async-safe storage for the current CUDA backend
+_CUDA_BACKEND: contextvars.ContextVar[str] = contextvars.ContextVar(
+    "_CUDA_BACKEND", default="ffmpeg"
+)
+
+
+@contextmanager
+def set_cuda_backend(backend: str) -> Generator[None, None, None]:
+    """Context Manager to set the CUDA backend for :class:`~torchcodec.decoders.VideoDecoder`.
+
+    This context manager allows you to specify which CUDA backend implementation
+    to use when creating :class:`~torchcodec.decoders.VideoDecoder` instances
+    with CUDA devices.
+
+    .. note::
+        **We recommend trying the "beta" backend instead of the default "ffmpeg"
+        backend!** The beta backend is faster, and will eventually become the
+        default in future versions. It may have rough edges that we'll polish
+        over time, but it's already quite stable and ready for adoption. Let us
+        know what you think!
+
+    Only the creation of the decoder needs to be inside the context manager, the
+    decoding methods can be called outside of it. You still need to pass
+    ``device="cuda"`` when creating the
+    :class:`~torchcodec.decoders.VideoDecoder` instance. If a CUDA device isn't
+    specified, this context manager will have no effect. See example below.
+
+    This is thread-safe and async-safe.
+
+    Args:
+        backend (str): The CUDA backend to use. Can be "ffmpeg" (default) or
+            "beta". We recommend trying "beta" as it's faster!
+
+    Example:
+        >>> with set_cuda_backend("beta"):
+        ...     decoder = VideoDecoder("video.mp4", device="cuda")
+        ...
+        ... # Only the decoder creation needs to be part of the context manager.
+        ... # Decoder will now the beta CUDA implementation:
+        ... decoder.get_frame_at(0)
+    """
+    backend = backend.lower()
+    if backend not in ("ffmpeg", "beta"):
+        raise ValueError(
+            f"Invalid CUDA backend ({backend}). Supported values are 'ffmpeg' and 'beta'."
+        )
+
+    previous_state = _CUDA_BACKEND.set(backend)
+    try:
+        yield
+    finally:
+        _CUDA_BACKEND.reset(previous_state)
+
+
+def _get_cuda_backend() -> str:
+    return _CUDA_BACKEND.get()