# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
from itertools import product
import random
import numpy as np
import torch
import torchaudio
from audiocraft.data.audio import audio_info, audio_read, audio_write, _av_read
from ..common_utils import TempDirMixin, get_white_noise, save_wav
class TestInfo(TempDirMixin):
def test_info_mp3(self):
sample_rates = [8000, 16_000]
channels = [1, 2]
duration = 1.
for sample_rate, ch in product(sample_rates, channels):
wav = get_white_noise(ch, int(sample_rate * duration))
path = self.get_temp_path('sample_wav.mp3')
save_wav(path, wav, sample_rate)
info = audio_info(path)
assert info.sample_rate == sample_rate
assert info.channels == ch
# we cannot trust torchaudio for num_frames, so we don't check
def _test_info_format(self, ext: str):
sample_rates = [8000, 16_000]
channels = [1, 2]
duration = 1.
for sample_rate, ch in product(sample_rates, channels):
n_frames = int(sample_rate * duration)
wav = get_white_noise(ch, n_frames)
path = self.get_temp_path(f'sample_wav{ext}')
save_wav(path, wav, sample_rate)
info = audio_info(path)
assert info.sample_rate == sample_rate
assert info.channels == ch
assert np.isclose(info.duration, duration, atol=1e-5)
def test_info_wav(self):
self._test_info_format('.wav')
def test_info_flac(self):
self._test_info_format('.flac')
def test_info_ogg(self):
self._test_info_format('.ogg')
def test_info_m4a(self):
# TODO: generate m4a file programmatically
# self._test_info_format('.m4a')
pass
class TestRead(TempDirMixin):
def test_read_full_wav(self):
sample_rates = [8000, 16_000]
channels = [1, 2]
duration = 1.
for sample_rate, ch in product(sample_rates, channels):
n_frames = int(sample_rate * duration)
wav = get_white_noise(ch, n_frames).clamp(-0.99, 0.99)
path = self.get_temp_path('sample_wav.wav')
save_wav(path, wav, sample_rate)
read_wav, read_sr = audio_read(path)
assert read_sr == sample_rate
assert read_wav.shape[0] == wav.shape[0]
assert read_wav.shape[1] == wav.shape[1]
assert torch.allclose(read_wav, wav, rtol=1e-03, atol=1e-04)
def test_read_partial_wav(self):
sample_rates = [8000, 16_000]
channels = [1, 2]
duration = 1.
read_duration = torch.rand(1).item()
for sample_rate, ch in product(sample_rates, channels):
n_frames = int(sample_rate * duration)
read_frames = int(sample_rate * read_duration)
wav = get_white_noise(ch, n_frames).clamp(-0.99, 0.99)
path = self.get_temp_path('sample_wav.wav')
save_wav(path, wav, sample_rate)
read_wav, read_sr = audio_read(path, 0, read_duration)
assert read_sr == sample_rate
assert read_wav.shape[0] == wav.shape[0]
assert read_wav.shape[1] == read_frames
assert torch.allclose(read_wav[..., 0:read_frames], wav[..., 0:read_frames], rtol=1e-03, atol=1e-04)
def test_read_seek_time_wav(self):
sample_rates = [8000, 16_000]
channels = [1, 2]
duration = 1.
read_duration = 1.
for sample_rate, ch in product(sample_rates, channels):
n_frames = int(sample_rate * duration)
wav = get_white_noise(ch, n_frames).clamp(-0.99, 0.99)
path = self.get_temp_path('sample_wav.wav')
save_wav(path, wav, sample_rate)
seek_time = torch.rand(1).item()
read_wav, read_sr = audio_read(path, seek_time, read_duration)
seek_frames = int(sample_rate * seek_time)
expected_frames = n_frames - seek_frames
assert read_sr == sample_rate
assert read_wav.shape[0] == wav.shape[0]
assert read_wav.shape[1] == expected_frames
assert torch.allclose(read_wav, wav[..., seek_frames:], rtol=1e-03, atol=1e-04)
def test_read_seek_time_wav_padded(self):
sample_rates = [8000, 16_000]
channels = [1, 2]
duration = 1.
read_duration = 1.
for sample_rate, ch in product(sample_rates, channels):
n_frames = int(sample_rate * duration)
read_frames = int(sample_rate * read_duration)
wav = get_white_noise(ch, n_frames).clamp(-0.99, 0.99)
path = self.get_temp_path('sample_wav.wav')
save_wav(path, wav, sample_rate)
seek_time = torch.rand(1).item()
seek_frames = int(sample_rate * seek_time)
expected_frames = n_frames - seek_frames
read_wav, read_sr = audio_read(path, seek_time, read_duration, pad=True)
expected_pad_wav = torch.zeros(wav.shape[0], read_frames - expected_frames)
assert read_sr == sample_rate
assert read_wav.shape[0] == wav.shape[0]
assert read_wav.shape[1] == read_frames
assert torch.allclose(read_wav[..., :expected_frames], wav[..., seek_frames:], rtol=1e-03, atol=1e-04)
assert torch.allclose(read_wav[..., expected_frames:], expected_pad_wav)
class TestAvRead(TempDirMixin):
def test_avread_seek_base(self):
sample_rates = [8000, 16_000]
channels = [1, 2]
duration = 2.
for sample_rate, ch in product(sample_rates, channels):
n_frames = int(sample_rate * duration)
wav = get_white_noise(ch, n_frames)
path = self.get_temp_path(f'reference_a_{sample_rate}_{ch}.wav')
save_wav(path, wav, sample_rate)
for _ in range(100):
# seek will always load a full duration segment in the file
seek_time = random.uniform(0.0, 1.0)
seek_duration = random.uniform(0.001, 1.0)
read_wav, read_sr = _av_read(path, seek_time, seek_duration)
assert read_sr == sample_rate
assert read_wav.shape[0] == wav.shape[0]
assert read_wav.shape[-1] == int(seek_duration * sample_rate)
def test_avread_seek_partial(self):
sample_rates = [8000, 16_000]
channels = [1, 2]
duration = 1.
for sample_rate, ch in product(sample_rates, channels):
n_frames = int(sample_rate * duration)
wav = get_white_noise(ch, n_frames)
path = self.get_temp_path(f'reference_b_{sample_rate}_{ch}.wav')
save_wav(path, wav, sample_rate)
for _ in range(100):
# seek will always load a partial segment
seek_time = random.uniform(0.5, 1.)
seek_duration = 1.
expected_num_frames = n_frames - int(seek_time * sample_rate)
read_wav, read_sr = _av_read(path, seek_time, seek_duration)
assert read_sr == sample_rate
assert read_wav.shape[0] == wav.shape[0]
assert read_wav.shape[-1] == expected_num_frames
def test_avread_seek_outofbound(self):
sample_rates = [8000, 16_000]
channels = [1, 2]
duration = 1.
for sample_rate, ch in product(sample_rates, channels):
n_frames = int(sample_rate * duration)
wav = get_white_noise(ch, n_frames)
path = self.get_temp_path(f'reference_c_{sample_rate}_{ch}.wav')
save_wav(path, wav, sample_rate)
seek_time = 1.5
read_wav, read_sr = _av_read(path, seek_time, 1.)
assert read_sr == sample_rate
assert read_wav.shape[0] == wav.shape[0]
assert read_wav.shape[-1] == 0
def test_avread_seek_edge(self):
sample_rates = [8000, 16_000]
# some of these values will have
# int(((frames - 1) / sample_rate) * sample_rate) != (frames - 1)
n_frames = [1000, 1001, 1002]
channels = [1, 2]
for sample_rate, ch, frames in product(sample_rates, channels, n_frames):
duration = frames / sample_rate
wav = get_white_noise(ch, frames)
path = self.get_temp_path(f'reference_d_{sample_rate}_{ch}.wav')
save_wav(path, wav, sample_rate)
seek_time = (frames - 1) / sample_rate
seek_frames = int(seek_time * sample_rate)
read_wav, read_sr = _av_read(path, seek_time, duration)
assert read_sr == sample_rate
assert read_wav.shape[0] == wav.shape[0]
assert read_wav.shape[-1] == (frames - seek_frames)
class TestAudioWrite(TempDirMixin):
def test_audio_write_wav(self):
torch.manual_seed(1234)
sample_rates = [8000, 16_000]
n_frames = [1000, 1001, 1002]
channels = [1, 2]
strategies = ["peak", "clip", "rms"]
formats = ["wav", "mp3"]
for sample_rate, ch, frames in product(sample_rates, channels, n_frames):
for format_, strategy in product(formats, strategies):
wav = get_white_noise(ch, frames)
path = self.get_temp_path(f'pred_{sample_rate}_{ch}')
audio_write(path, wav, sample_rate, format_, strategy=strategy)
read_wav, read_sr = torchaudio.load(f'{path}.{format_}')
if format_ == "wav":
assert read_wav.shape == wav.shape
if format_ == "wav" and strategy in ["peak", "rms"]:
rescaled_read_wav = read_wav / read_wav.abs().max() * wav.abs().max()
# for a Gaussian, the typical max scale will be less than ~5x the std.
# The error when writing to disk will ~ 1/2**15, and when rescaling, 5x that.
# For RMS target, rescaling leaves more headroom by default, leading
# to a 20x rescaling typically
atol = (5 if strategy == "peak" else 20) / 2**15
delta = (rescaled_read_wav - wav).abs().max()
assert torch.allclose(wav, rescaled_read_wav, rtol=0, atol=atol), (delta, atol)
formats = ["wav"] # faster unit tests