audiocraft/tests/data/test_audio.py

# 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