test_datasets.py 3.06 KB
Newer Older
1
import os
2
3
4
5
import unittest

from torchaudio.datasets.commonvoice import COMMONVOICE
from torchaudio.datasets.librispeech import LIBRISPEECH
6
from torchaudio.datasets.speechcommands import SPEECHCOMMANDS
Vincent QB's avatar
Vincent QB committed
7
from torchaudio.datasets.utils import diskcache_iterator, bg_iterator
8
9
from torchaudio.datasets.vctk import VCTK
from torchaudio.datasets.yesno import YESNO
Taras Sereda's avatar
Taras Sereda committed
10
from torchaudio.datasets.ljspeech import LJSPEECH
11
from torchaudio.datasets.gtzan import GTZAN
jimchen90's avatar
jimchen90 committed
12
from torchaudio.datasets.cmuarctic import CMUARCTIC
13

14
15
16
17
18
19
20
21
from .common_utils import (
    TempDirMixin,
    TorchaudioTestCase,
    get_asset_path,
    get_whitenoise,
    save_wav,
    normalize_wav,
)
Vincent QB's avatar
Vincent QB committed
22

23

24
class TestDatasets(TorchaudioTestCase):
moto's avatar
moto committed
25
    backend = 'default'
26
    path = get_asset_path()
27
28

    def test_vctk(self):
Vincent QB's avatar
Vincent QB committed
29
        data = VCTK(self.path)
30
31
32
33
34
35
        data[0]

    def test_librispeech(self):
        data = LIBRISPEECH(self.path, "dev-clean")
        data[0]

moto's avatar
moto committed
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
    def test_ljspeech(self):
        data = LJSPEECH(self.path)
        data[0]

    def test_speechcommands(self):
        data = SPEECHCOMMANDS(self.path)
        data[0]

    def test_gtzan(self):
        data = GTZAN(self.path)
        data[0]

    def test_cmuarctic(self):
        data = CMUARCTIC(self.path)
        data[0]


53
class TestCommonVoice(TorchaudioTestCase):
54
    backend = 'default'
55
    path = get_asset_path()
moto's avatar
moto committed
56

57
    def test_commonvoice(self):
58
        data = COMMONVOICE(self.path, url="tatar")
59
60
61
        data[0]

    def test_commonvoice_diskcache(self):
62
        data = COMMONVOICE(self.path, url="tatar")
Vincent QB's avatar
Vincent QB committed
63
        data = diskcache_iterator(data)
64
65
66
67
68
        # Save
        data[0]
        # Load
        data[0]

Vincent QB's avatar
Vincent QB committed
69
    def test_commonvoice_bg(self):
70
        data = COMMONVOICE(self.path, url="tatar")
Vincent QB's avatar
Vincent QB committed
71
        data = bg_iterator(data, 5)
moto's avatar
moto committed
72
        for _ in data:
Vincent QB's avatar
Vincent QB committed
73
74
            pass

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
class TestYesNo(TempDirMixin, TorchaudioTestCase):
    backend = 'default'

    root_dir = None
    data = []
    labels = [
        [0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 1, 1, 1, 1],
        [0, 1, 0, 1, 0, 1, 1, 0],
        [1, 1, 1, 1, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 1, 1, 1],
    ]

    @classmethod
    def setUpClass(cls):
        cls.root_dir = cls.get_base_temp_dir()
        base_dir = os.path.join(cls.root_dir, 'waves_yesno')
        os.makedirs(base_dir, exist_ok=True)
        for label in cls.labels:
            filename = f'{"_".join(str(l) for l in label)}.wav'
            path = os.path.join(base_dir, filename)
            data = get_whitenoise(sample_rate=8000, duration=6, n_channels=1, dtype='int16')
            save_wav(path, data, 8000)
            cls.data.append(normalize_wav(data))

    def test_yesno(self):
        dataset = YESNO(self.root_dir)
        samples = list(dataset)
        samples.sort(key=lambda s: s[2])
        for i, (waveform, sample_rate, label) in enumerate(samples):
            expected_label = self.labels[i]
            expected_data = self.data[i]
            self.assertEqual(expected_data, waveform, atol=5e-5, rtol=1e-8)
            assert sample_rate == 8000
            assert label == expected_label


113
114
if __name__ == "__main__":
    unittest.main()