File size: 12,217 Bytes
864affd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
import io
from typing import Iterator, List, Optional

import torch
from torch import Tensor

from torio.io._streaming_media_decoder import _get_afilter_desc, StreamingMediaDecoder as StreamReader
from torio.io._streaming_media_encoder import CodecConfig, StreamingMediaEncoder as 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)