torchaudio 2.0.2__cp310-cp310-manylinux2014_aarch64.whl → 2.1.1__cp310-cp310-manylinux2014_aarch64.whl

torchaudio/io/_effector.py

@@ -0,0 +1,347 @@
+import io
+from typing import Iterator, List, Optional
+
+import torch
+from torch import Tensor
+
+from ._stream_reader import _get_afilter_desc, StreamReader
+from ._stream_writer import CodecConfig, StreamWriter
+
+
+class _StreamingIOBuffer:
+    """Streaming Bytes IO buffer. Data are dropped when read."""
+
+    def __init__(self):
+        self._buffer: List(bytes) = []
+
+    def write(self, b: bytes):
+        if b:
+            self._buffer.append(b)
+        return len(b)
+
+    def pop(self, n):
+        """Pop the oldest byte string. It does not necessary return the requested amount"""
+        if not self._buffer:
+            return b""
+        if len(self._buffer[0]) <= n:
+            return self._buffer.pop(0)
+        ret = self._buffer[0][:n]
+        self._buffer[0] = self._buffer[0][n:]
+        return ret
+
+
+def _get_sample_fmt(dtype: torch.dtype):
+    types = {
+        torch.uint8: "u8",
+        torch.int16: "s16",
+        torch.int32: "s32",
+        torch.float32: "flt",
+        torch.float64: "dbl",
+    }
+    if dtype not in types:
+        raise ValueError(f"Unsupported dtype is provided {dtype}. Supported dtypes are: {types.keys()}")
+    return types[dtype]
+
+
+class _AudioStreamingEncoder:
+    """Given a waveform, encode on-demand and return bytes"""
+
+    def __init__(
+        self,
+        src: Tensor,
+        sample_rate: int,
+        effect: str,
+        muxer: str,
+        encoder: Optional[str],
+        codec_config: Optional[CodecConfig],
+        frames_per_chunk: int,
+    ):
+        self.src = src
+        self.buffer = _StreamingIOBuffer()
+        self.writer = StreamWriter(self.buffer, format=muxer)
+        self.writer.add_audio_stream(
+            num_channels=src.size(1),
+            sample_rate=sample_rate,
+            format=_get_sample_fmt(src.dtype),
+            encoder=encoder,
+            filter_desc=effect,
+            codec_config=codec_config,
+        )
+        self.writer.open()
+        self.fpc = frames_per_chunk
+
+        # index on the input tensor (along time-axis)
+        # we use -1 to indicate that we finished iterating the tensor and
+        # the writer is closed.
+        self.i_iter = 0
+
+    def read(self, n):
+        while not self.buffer._buffer and self.i_iter >= 0:
+            self.writer.write_audio_chunk(0, self.src[self.i_iter : self.i_iter + self.fpc])
+            self.i_iter += self.fpc
+            if self.i_iter >= self.src.size(0):
+                self.writer.flush()
+                self.writer.close()
+                self.i_iter = -1
+        return self.buffer.pop(n)
+
+
+def _encode(
+    src: Tensor,
+    sample_rate: int,
+    effect: str,
+    muxer: str,
+    encoder: Optional[str],
+    codec_config: Optional[CodecConfig],
+):
+    buffer = io.BytesIO()
+    writer = StreamWriter(buffer, format=muxer)
+    writer.add_audio_stream(
+        num_channels=src.size(1),
+        sample_rate=sample_rate,
+        format=_get_sample_fmt(src.dtype),
+        encoder=encoder,
+        filter_desc=effect,
+        codec_config=codec_config,
+    )
+    with writer.open():
+        writer.write_audio_chunk(0, src)
+    buffer.seek(0)
+    return buffer
+
+
+def _get_muxer(dtype: torch.dtype):
+    # TODO: check if this works in Windows.
+    types = {
+        torch.uint8: "u8",
+        torch.int16: "s16le",
+        torch.int32: "s32le",
+        torch.float32: "f32le",
+        torch.float64: "f64le",
+    }
+    if dtype not in types:
+        raise ValueError(f"Unsupported dtype is provided {dtype}. Supported dtypes are: {types.keys()}")
+    return types[dtype]
+
+
+class AudioEffector:
+    """Apply various filters and/or codecs to waveforms.
+
+    .. versionadded:: 2.1
+
+    Args:
+        effect (str or None, optional): Filter expressions or ``None`` to apply no filter.
+            See https://ffmpeg.org/ffmpeg-filters.html#Audio-Filters for the
+            details of filter syntax.
+
+        format (str or None, optional): When provided, encode the audio into the
+            corresponding format. Default: ``None``.
+
+        encoder (str or None, optional): When provided, override the encoder used
+            by the ``format``. Default: ``None``.
+
+        codec_config (CodecConfig or None, optional): When provided, configure the encoding codec.
+            Should be provided in conjunction with ``format`` option.
+
+        pad_end (bool, optional): When enabled, and if the waveform becomes shorter after applying
+            effects/codec, then pad the end with silence.
+
+    Example - Basic usage
+        To use ``AudioEffector``, first instantiate it with a set of
+        ``effect`` and ``format``.
+
+        >>> # instantiate the effector
+        >>> effector = AudioEffector(effect=..., format=...)
+
+        Then, use :py:meth:`~AudioEffector.apply` or :py:meth:`~AudioEffector.stream`
+        method to apply them.
+
+        >>> # Apply the effect to the whole waveform
+        >>> applied = effector.apply(waveform, sample_rate)
+
+        >>> # Apply the effect chunk-by-chunk
+        >>> for chunk in effector.stream(waveform, sample_rate):
+        >>>    ...
+
+    Example - Applying effects
+        Please refer to
+        https://ffmpeg.org/ffmpeg-filters.html#Filtergraph-description
+        for the overview of filter description, and
+        https://ffmpeg.org/ffmpeg-filters.html#toc-Audio-Filters
+        for the list of available filters.
+
+        Tempo - https://ffmpeg.org/ffmpeg-filters.html#atempo
+
+        >>> AudioEffector(effect="atempo=1.5")
+
+        Echo - https://ffmpeg.org/ffmpeg-filters.html#aecho
+
+        >>> AudioEffector(effect="aecho=0.8:0.88:60:0.4")
+
+        Flanger - https://ffmpeg.org/ffmpeg-filters.html#flanger
+
+        >>> AudioEffector(effect="aflanger")
+
+        Vibrato - https://ffmpeg.org/ffmpeg-filters.html#vibrato
+
+        >>> AudioEffector(effect="vibrato")
+
+        Tremolo - https://ffmpeg.org/ffmpeg-filters.html#tremolo
+
+        >>> AudioEffector(effect="vibrato")
+
+        You can also apply multiple effects at once.
+
+        >>> AudioEffector(effect="")
+
+    Example - Applying codec
+        One can apply codec using ``format`` argument. ``format`` can be
+        audio format or container format. If the container format supports
+        multiple encoders, you can specify it with ``encoder`` argument.
+
+        Wav format
+        (no compression is applied but samples are converted to
+        16-bit signed integer)
+
+        >>> AudioEffector(format="wav")
+
+        Ogg format with default encoder
+
+        >>> AudioEffector(format="ogg")
+
+        Ogg format with vorbis
+
+        >>> AudioEffector(format="ogg", encoder="vorbis")
+
+        Ogg format with opus
+
+        >>> AudioEffector(format="ogg", encoder="opus")
+
+        Webm format with opus
+
+        >>> AudioEffector(format="webm", encoder="opus")
+
+    Example - Applying codec with configuration
+        Reference: https://trac.ffmpeg.org/wiki/Encode/MP3
+
+        MP3 with default config
+
+        >>> AudioEffector(format="mp3")
+
+        MP3 with variable bitrate
+
+        >>> AudioEffector(format="mp3", codec_config=CodecConfig(qscale=5))
+
+        MP3 with constant bitrate
+
+        >>> AudioEffector(format="mp3", codec_config=CodecConfig(bit_rate=32_000))
+    """
+
+    def __init__(
+        self,
+        effect: Optional[str] = None,
+        format: Optional[str] = None,
+        *,
+        encoder: Optional[str] = None,
+        codec_config: Optional[CodecConfig] = None,
+        pad_end: bool = True,
+    ):
+        if format is None:
+            if encoder is not None or codec_config is not None:
+                raise ValueError("`encoder` and/or `condec_config` opions are provided without `format` option.")
+        self.effect = effect
+        self.format = format
+        self.encoder = encoder
+        self.codec_config = codec_config
+        self.pad_end = pad_end
+
+    def _get_reader(self, waveform, sample_rate, output_sample_rate, frames_per_chunk=None):
+        num_frames, num_channels = waveform.shape
+
+        if self.format is not None:
+            muxer = self.format
+            encoder = self.encoder
+            option = {}
+            # Some formats are headerless, so need to provide these infomation.
+            if self.format == "mulaw":
+                option = {"sample_rate": f"{sample_rate}", "channels": f"{num_channels}"}
+
+        else:  # PCM
+            muxer = _get_muxer(waveform.dtype)
+            encoder = None
+            option = {"sample_rate": f"{sample_rate}", "channels": f"{num_channels}"}
+
+        if frames_per_chunk is None:
+            src = _encode(waveform, sample_rate, self.effect, muxer, encoder, self.codec_config)
+        else:
+            src = _AudioStreamingEncoder(
+                waveform, sample_rate, self.effect, muxer, encoder, self.codec_config, frames_per_chunk
+            )
+
+        output_sr = sample_rate if output_sample_rate is None else output_sample_rate
+        filter_desc = _get_afilter_desc(output_sr, _get_sample_fmt(waveform.dtype), num_channels)
+        if self.pad_end:
+            filter_desc = f"{filter_desc},apad=whole_len={num_frames}"
+
+        reader = StreamReader(src, format=muxer, option=option)
+        reader.add_audio_stream(frames_per_chunk or -1, -1, filter_desc=filter_desc)
+        return reader
+
+    def apply(self, waveform: Tensor, sample_rate: int, output_sample_rate: Optional[int] = None) -> Tensor:
+        """Apply the effect and/or codecs to the whole tensor.
+
+        Args:
+            waveform (Tensor): The input waveform. Shape: ``(time, channel)``
+            sample_rate (int): Sample rate of the input waveform.
+            output_sample_rate (int or None, optional): Output sample rate.
+                If provided, override the output sample rate.
+                Otherwise, the resulting tensor is resampled to have
+                the same sample rate as the input.
+                Default: ``None``.
+
+        Returns:
+            Tensor:
+                Resulting Tensor. Shape: ``(time, channel)``. The number of frames
+                could be different from that of the input.
+        """
+        if waveform.ndim != 2:
+            raise ValueError(f"Expected the input waveform to be 2D. Found: {waveform.ndim}")
+
+        if waveform.numel() == 0:
+            return waveform
+
+        reader = self._get_reader(waveform, sample_rate, output_sample_rate)
+        reader.process_all_packets()
+        (applied,) = reader.pop_chunks()
+        return Tensor(applied)
+
+    def stream(
+        self, waveform: Tensor, sample_rate: int, frames_per_chunk: int, output_sample_rate: Optional[int] = None
+    ) -> Iterator[Tensor]:
+        """Apply the effect and/or codecs to the given tensor chunk by chunk.
+
+        Args:
+            waveform (Tensor): The input waveform. Shape: ``(time, channel)``
+            sample_rate (int): Sample rate of the waveform.
+            frames_per_chunk (int): The number of frames to return at a time.
+            output_sample_rate (int or None, optional): Output sample rate.
+                If provided, override the output sample rate.
+                Otherwise, the resulting tensor is resampled to have
+                the same sample rate as the input.
+                Default: ``None``.
+
+        Returns:
+            Iterator[Tensor]:
+                Series of processed chunks. Shape: ``(time, channel)``, where the
+                the number of frames matches ``frames_per_chunk`` except the
+                last chunk, which could be shorter.
+        """
+        if waveform.ndim != 2:
+            raise ValueError(f"Expected the input waveform to be 2D. Found: {waveform.ndim}")
+
+        if waveform.numel() == 0:
+            return waveform
+
+        reader = self._get_reader(waveform, sample_rate, output_sample_rate, frames_per_chunk)
+        for (applied,) in reader.stream():
+            yield Tensor(applied)