test_dataloader.py

import unittest

import torchaudio
from torch.utils.data import Dataset, DataLoader

import common_utils
from common_utils import AudioBackendScope, BACKENDS


@unittest.skipIf("sox" not in BACKENDS, "sox not available")
class TORCHAUDIODS(Dataset):

    def __init__(self):
        sound_files = ["sinewave.wav", "steam-train-whistle-daniel_simon.mp3"]
        self.data = [common_utils.get_asset_path(fn) for fn in sound_files]
        self.si, self.ei = torchaudio.info(common_utils.get_asset_path("sinewave.wav"))
        self.si.precision = 16
        self.E = torchaudio.sox_effects.SoxEffectsChain()
        self.E.append_effect_to_chain("rate", [self.si.rate])  # resample to 16000hz
        self.E.append_effect_to_chain("channels", [self.si.channels])  # mono signal
        self.E.append_effect_to_chain("trim", [0, "16000s"])  # first 16000 samples of audio

    def __getitem__(self, index):
        fn = self.data[index]
        self.E.set_input_file(fn)
        x, sr = self.E.sox_build_flow_effects()
        return x

    def __len__(self):
        return len(self.data)


@unittest.skipIf("sox" not in BACKENDS, "sox not available")
class Test_DataLoader(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        torchaudio.initialize_sox()

    @classmethod
    def tearDownClass(cls):
        torchaudio.shutdown_sox()

    def test_1(self):
        expected_size = (2, 1, 16000)
        ds = TORCHAUDIODS()
        dl = DataLoader(ds, batch_size=2)
        for x in dl:
            self.assertTrue(x.size() == expected_size)

if __name__ == '__main__':
    with AudioBackendScope("sox"):
        unittest.main()