test_io.py 10.1 KB
Newer Older
moto's avatar
moto committed
1
2
import os
import math
David Pollack's avatar
David Pollack committed
3
import unittest
moto's avatar
moto committed
4

5
import torch
Soumith Chintala's avatar
Soumith Chintala committed
6
import torchaudio
moto's avatar
moto committed
7
8

from .common_utils import BACKENDS, BACKENDS_MP3, create_temp_assets_dir
Vincent QB's avatar
Vincent QB committed
9
10


David Pollack's avatar
David Pollack committed
11
class Test_LoadSave(unittest.TestCase):
12
    test_dirpath, test_dir = create_temp_assets_dir()
13
14
    test_filepath = os.path.join(test_dirpath, "assets",
                                 "steam-train-whistle-daniel_simon.mp3")
Vincent QB's avatar
Vincent QB committed
15
16
    test_filepath_wav = os.path.join(test_dirpath, "assets",
                                     "steam-train-whistle-daniel_simon.wav")
Soumith Chintala's avatar
Soumith Chintala committed
17

David Pollack's avatar
David Pollack committed
18
    def test_1_save(self):
19
        for backend in BACKENDS_MP3:
Vincent QB's avatar
Vincent QB committed
20
            with self.subTest():
moto's avatar
moto committed
21
22
                torchaudio.set_audio_backend(backend)
                self._test_1_save(self.test_filepath, False)
Vincent QB's avatar
Vincent QB committed
23

24
        for backend in BACKENDS:
Vincent QB's avatar
Vincent QB committed
25
            with self.subTest():
moto's avatar
moto committed
26
27
                torchaudio.set_audio_backend(backend)
                self._test_1_save(self.test_filepath_wav, True)
Vincent QB's avatar
Vincent QB committed
28
29

    def _test_1_save(self, test_filepath, normalization):
David Pollack's avatar
David Pollack committed
30
        # load signal
Vincent QB's avatar
Vincent QB committed
31
        x, sr = torchaudio.load(test_filepath, normalization=normalization)
David Pollack's avatar
David Pollack committed
32
33
34
35
36
37
38
39
40
41
42
43
44
45

        # check save
        new_filepath = os.path.join(self.test_dirpath, "test.wav")
        torchaudio.save(new_filepath, x, sr)
        self.assertTrue(os.path.isfile(new_filepath))
        os.unlink(new_filepath)

        # check automatic normalization
        x /= 1 << 31
        torchaudio.save(new_filepath, x, sr)
        self.assertTrue(os.path.isfile(new_filepath))
        os.unlink(new_filepath)

        # test save 1d tensor
David Pollack's avatar
David Pollack committed
46
        x = x[0, :]  # get mono signal
Soumith Chintala's avatar
Soumith Chintala committed
47
        x.squeeze_()  # remove channel dim
David Pollack's avatar
David Pollack committed
48
49
50
51
52
53
        torchaudio.save(new_filepath, x, sr)
        self.assertTrue(os.path.isfile(new_filepath))
        os.unlink(new_filepath)

        # don't allow invalid sizes as inputs
        with self.assertRaises(ValueError):
David Pollack's avatar
David Pollack committed
54
            x.unsqueeze_(1)  # L x C not C x L
David Pollack's avatar
David Pollack committed
55
56
57
58
59
            torchaudio.save(new_filepath, x, sr)

        with self.assertRaises(ValueError):
            x.squeeze_()
            x.unsqueeze_(1)
Soumith Chintala's avatar
Soumith Chintala committed
60
            x.unsqueeze_(0)  # 1 x L x 1
David Pollack's avatar
David Pollack committed
61
62
63
64
            torchaudio.save(new_filepath, x, sr)

        # don't save to folders that don't exist
        with self.assertRaises(OSError):
65
66
            new_filepath = os.path.join(self.test_dirpath, "no-path",
                                        "test.wav")
David Pollack's avatar
David Pollack committed
67
            torchaudio.save(new_filepath, x, sr)
Soumith Chintala's avatar
Soumith Chintala committed
68

Vincent QB's avatar
Vincent QB committed
69
    def test_1_save_sine(self):
70
        for backend in BACKENDS:
Vincent QB's avatar
Vincent QB committed
71
            with self.subTest():
moto's avatar
moto committed
72
73
                torchaudio.set_audio_backend(backend)
                self._test_1_save_sine()
Vincent QB's avatar
Vincent QB committed
74
75
76

    def _test_1_save_sine(self):

77
        # save created file
78
79
        sinewave_filepath = os.path.join(self.test_dirpath, "assets",
                                         "sinewave.wav")
80
81
82
83
        sr = 16000
        freq = 440
        volume = 0.3

84
        y = (torch.cos(
85
            2 * math.pi * torch.arange(0, 4 * sr).float() * freq / sr))
David Pollack's avatar
David Pollack committed
86
        y.unsqueeze_(0)
87
        # y is between -1 and 1, so must scale
David Pollack's avatar
David Pollack committed
88
        y = (y * volume * (2**31)).long()
89
90
        torchaudio.save(sinewave_filepath, y, sr)
        self.assertTrue(os.path.isfile(sinewave_filepath))
91

92
        # test precision
David Pollack's avatar
David Pollack committed
93
        new_precision = 32
94
        new_filepath = os.path.join(self.test_dirpath, "test.wav")
David Pollack's avatar
David Pollack committed
95
96
97
98
99
        si, ei = torchaudio.info(sinewave_filepath)
        torchaudio.save(new_filepath, y, sr, new_precision)
        si32, ei32 = torchaudio.info(new_filepath)
        self.assertEqual(si.precision, 16)
        self.assertEqual(si32.precision, new_precision)
100
101
        os.unlink(new_filepath)

David Pollack's avatar
David Pollack committed
102
    def test_2_load(self):
103
        for backend in BACKENDS_MP3:
Vincent QB's avatar
Vincent QB committed
104
            with self.subTest():
moto's avatar
moto committed
105
106
                torchaudio.set_audio_backend(backend)
                self._test_2_load(self.test_filepath, 278756)
Vincent QB's avatar
Vincent QB committed
107

108
        for backend in BACKENDS:
Vincent QB's avatar
Vincent QB committed
109
            with self.subTest():
moto's avatar
moto committed
110
111
                torchaudio.set_audio_backend(backend)
                self._test_2_load(self.test_filepath_wav, 276858)
Vincent QB's avatar
Vincent QB committed
112
113

    def _test_2_load(self, test_filepath, length):
David Pollack's avatar
David Pollack committed
114
        # check normal loading
Vincent QB's avatar
Vincent QB committed
115
        x, sr = torchaudio.load(test_filepath)
David Pollack's avatar
David Pollack committed
116
        self.assertEqual(sr, 44100)
Vincent QB's avatar
Vincent QB committed
117
        self.assertEqual(x.size(), (2, length))
David Pollack's avatar
David Pollack committed
118
119
120

        # check offset
        offset = 15
Vincent QB's avatar
Vincent QB committed
121
122
        x, _ = torchaudio.load(test_filepath)
        x_offset, _ = torchaudio.load(test_filepath, offset=offset)
123
        self.assertTrue(x[:, offset:].allclose(x_offset))
David Pollack's avatar
David Pollack committed
124
125
126

        # check number of frames
        n = 201
Vincent QB's avatar
Vincent QB committed
127
        x, _ = torchaudio.load(test_filepath, num_frames=n)
David Pollack's avatar
David Pollack committed
128
129
130
        self.assertTrue(x.size(), (2, n))

        # check channels first
Vincent QB's avatar
Vincent QB committed
131
132
        x, _ = torchaudio.load(test_filepath, channels_first=False)
        self.assertEqual(x.size(), (length, 2))
David Pollack's avatar
David Pollack committed
133
134
135
136
137
138
139
140
141
142

        # check raising errors
        with self.assertRaises(OSError):
            torchaudio.load("file-does-not-exist.mp3")

        with self.assertRaises(OSError):
            tdir = os.path.join(
                os.path.dirname(self.test_dirpath), "torchaudio")
            torchaudio.load(tdir)

Vincent QB's avatar
Vincent QB committed
143
    def test_2_load_nonormalization(self):
144
        for backend in BACKENDS_MP3:
Vincent QB's avatar
Vincent QB committed
145
            with self.subTest():
moto's avatar
moto committed
146
147
                torchaudio.set_audio_backend(backend)
                self._test_2_load_nonormalization(self.test_filepath, 278756)
Vincent QB's avatar
Vincent QB committed
148
149
150
151
152
153
154
155
156
157
158
159

    def _test_2_load_nonormalization(self, test_filepath, length):

        # check no normalizing
        x, _ = torchaudio.load(test_filepath, normalization=False)
        self.assertTrue(x.min() <= -1.0)
        self.assertTrue(x.max() >= 1.0)

        # check different input tensor type
        x, _ = torchaudio.load(test_filepath, torch.LongTensor(), normalization=False)
        self.assertTrue(isinstance(x, torch.LongTensor))

David Pollack's avatar
David Pollack committed
160
    def test_3_load_and_save_is_identity(self):
161
        for backend in BACKENDS:
Vincent QB's avatar
Vincent QB committed
162
            with self.subTest():
moto's avatar
moto committed
163
164
                torchaudio.set_audio_backend(backend)
                self._test_3_load_and_save_is_identity()
Vincent QB's avatar
Vincent QB committed
165
166

    def _test_3_load_and_save_is_identity(self):
167
168
169
170
171
172
173
174
175
        input_path = os.path.join(self.test_dirpath, 'assets', 'sinewave.wav')
        tensor, sample_rate = torchaudio.load(input_path)
        output_path = os.path.join(self.test_dirpath, 'test.wav')
        torchaudio.save(output_path, tensor, sample_rate)
        tensor2, sample_rate2 = torchaudio.load(output_path)
        self.assertTrue(tensor.allclose(tensor2))
        self.assertEqual(sample_rate, sample_rate2)
        os.unlink(output_path)

moto's avatar
moto committed
176
    @unittest.skipIf(any(be not in BACKENDS for be in ["sox", "soundfile"]), "sox and soundfile are not available")
Vincent QB's avatar
Vincent QB committed
177
178
179
180
181
182
183
    def test_3_load_and_save_is_identity_across_backend(self):
        with self.subTest():
            self._test_3_load_and_save_is_identity_across_backend("sox", "soundfile")
        with self.subTest():
            self._test_3_load_and_save_is_identity_across_backend("soundfile", "sox")

    def _test_3_load_and_save_is_identity_across_backend(self, backend1, backend2):
moto's avatar
moto committed
184
185
186
        torchaudio.set_audio_backend(backend1)
        input_path = os.path.join(self.test_dirpath, 'assets', 'sinewave.wav')
        tensor1, sample_rate1 = torchaudio.load(input_path)
Vincent QB's avatar
Vincent QB committed
187

moto's avatar
moto committed
188
189
        output_path = os.path.join(self.test_dirpath, 'test.wav')
        torchaudio.save(output_path, tensor1, sample_rate1)
Vincent QB's avatar
Vincent QB committed
190

moto's avatar
moto committed
191
192
        torchaudio.set_audio_backend(backend2)
        tensor2, sample_rate2 = torchaudio.load(output_path)
Vincent QB's avatar
Vincent QB committed
193
194
195
196
197

        self.assertTrue(tensor1.allclose(tensor2))
        self.assertEqual(sample_rate1, sample_rate2)
        os.unlink(output_path)

David Pollack's avatar
David Pollack committed
198
    def test_4_load_partial(self):
199
        for backend in BACKENDS_MP3:
Vincent QB's avatar
Vincent QB committed
200
            with self.subTest():
moto's avatar
moto committed
201
202
                torchaudio.set_audio_backend(backend)
                self._test_4_load_partial()
Vincent QB's avatar
Vincent QB committed
203
204

    def _test_4_load_partial(self):
David Pollack's avatar
David Pollack committed
205
206
        num_frames = 101
        offset = 201
207
208
209
210
        # load entire mono sinewave wav file, load a partial copy and then compare
        input_sine_path = os.path.join(self.test_dirpath, 'assets', 'sinewave.wav')
        x_sine_full, sr_sine = torchaudio.load(input_sine_path)
        x_sine_part, _ = torchaudio.load(input_sine_path, num_frames=num_frames, offset=offset)
211
        l1_error = x_sine_full[:, offset:(num_frames + offset)].sub(x_sine_part).abs().sum().item()
212
        # test for the correct number of samples and that the correct portion was loaded
David Pollack's avatar
David Pollack committed
213
        self.assertEqual(x_sine_part.size(1), num_frames)
214
215
216
217
218
219
220
        self.assertEqual(l1_error, 0.)
        # create a two channel version of this wavefile
        x_2ch_sine = x_sine_full.repeat(1, 2)
        out_2ch_sine_path = os.path.join(self.test_dirpath, 'assets', '2ch_sinewave.wav')
        torchaudio.save(out_2ch_sine_path, x_2ch_sine, sr_sine)
        x_2ch_sine_load, _ = torchaudio.load(out_2ch_sine_path, num_frames=num_frames, offset=offset)
        os.unlink(out_2ch_sine_path)
David Pollack's avatar
David Pollack committed
221
        l1_error = x_2ch_sine_load.sub(x_2ch_sine[:, offset:(offset + num_frames)]).abs().sum().item()
222
223
224
225
226
        self.assertEqual(l1_error, 0.)

        # test with two channel mp3
        x_2ch_full, sr_2ch = torchaudio.load(self.test_filepath, normalization=True)
        x_2ch_part, _ = torchaudio.load(self.test_filepath, normalization=True, num_frames=num_frames, offset=offset)
227
        l1_error = x_2ch_full[:, offset:(offset + num_frames)].sub(x_2ch_part).abs().sum().item()
David Pollack's avatar
David Pollack committed
228
        self.assertEqual(x_2ch_part.size(1), num_frames)
229
230
231
232
233
        self.assertEqual(l1_error, 0.)

        # check behavior if number of samples would exceed file length
        offset_ns = 300
        x_ns, _ = torchaudio.load(input_sine_path, num_frames=100000, offset=offset_ns)
David Pollack's avatar
David Pollack committed
234
        self.assertEqual(x_ns.size(1), x_sine_full.size(1) - offset_ns)
235
236
237
238
239

        # check when offset is beyond the end of the file
        with self.assertRaises(RuntimeError):
            torchaudio.load(input_sine_path, offset=100000)

David Pollack's avatar
David Pollack committed
240
    def test_5_get_info(self):
241
        for backend in BACKENDS:
Vincent QB's avatar
Vincent QB committed
242
            with self.subTest():
moto's avatar
moto committed
243
244
                torchaudio.set_audio_backend(backend)
                self._test_5_get_info()
Vincent QB's avatar
Vincent QB committed
245
246

    def _test_5_get_info(self):
247
        input_path = os.path.join(self.test_dirpath, 'assets', 'sinewave.wav')
David Pollack's avatar
David Pollack committed
248
249
250
251
252
253
        channels, samples, rate, precision = (1, 64000, 16000, 16)
        si, ei = torchaudio.info(input_path)
        self.assertEqual(si.channels, channels)
        self.assertEqual(si.length, samples)
        self.assertEqual(si.rate, rate)
        self.assertEqual(ei.bits_per_sample, precision)
Soumith Chintala's avatar
Soumith Chintala committed
254

David Pollack's avatar
David Pollack committed
255
256
if __name__ == '__main__':
    unittest.main()