Python torchaudio.save() Examples

def _create_data_set(self):
        # used to generate the dataset to test on. this is not used in testing (offline procedure)
        test_filepath = common_utils.get_asset_path('kaldi_file.wav')
        sr = 16000
        x = torch.arange(0, 20).float()
        # between [-6,6]
        y = torch.cos(2 * math.pi * x) + 3 * torch.sin(math.pi * x) + 2 * torch.cos(x)
        # between [-2^30, 2^30]
        y = (y / 6 * (1 << 30)).long()
        # clear the last 16 bits because they aren't used anyways
        y = ((y >> 16) << 16).float()
        torchaudio.save(test_filepath, y, sr)
        sound, sample_rate = torchaudio.load(test_filepath, normalization=False)
        print(y >> 16)
        self.assertTrue(sample_rate == sr)
        torch.testing.assert_allclose(y, sound) 

def test_load_wav(self, dtype, sample_rate, num_channels, normalize, channels_first):
        """`sox_io_backend.load` is torchscript-able and returns the same result"""
        audio_path = self.get_temp_path(f'test_load_{dtype}_{sample_rate}_{num_channels}_{normalize}.wav')
        data = get_wav_data(dtype, num_channels, normalize=False, num_frames=1 * sample_rate)
        save_wav(audio_path, data, sample_rate)

        script_path = self.get_temp_path('load_func.zip')
        torch.jit.script(py_load_func).save(script_path)
        ts_load_func = torch.jit.load(script_path)

        py_data, py_sr = py_load_func(
            audio_path, normalize=normalize, channels_first=channels_first)
        ts_data, ts_sr = ts_load_func(
            audio_path, normalize=normalize, channels_first=channels_first)

        self.assertEqual(py_sr, ts_sr)
        self.assertEqual(py_data, ts_data) 

def test_save_wav(self, dtype, sample_rate, num_channels):
        script_path = self.get_temp_path('save_func.zip')
        torch.jit.script(py_save_func).save(script_path)
        ts_save_func = torch.jit.load(script_path)

        expected = get_wav_data(dtype, num_channels)
        py_path = self.get_temp_path(f'test_save_py_{dtype}_{sample_rate}_{num_channels}.wav')
        ts_path = self.get_temp_path(f'test_save_ts_{dtype}_{sample_rate}_{num_channels}.wav')

        py_save_func(py_path, expected, sample_rate, True, None)
        ts_save_func(ts_path, expected, sample_rate, True, None)

        py_data, py_sr = load_wav(py_path)
        ts_data, ts_sr = load_wav(ts_path)

        self.assertEqual(sample_rate, py_sr)
        self.assertEqual(sample_rate, ts_sr)
        self.assertEqual(expected, py_data)
        self.assertEqual(expected, ts_data) 

def setUpClass(cls):
        if not os.path.exists(cls._AUDIO_DATA_DIR):
            os.makedirs(cls._AUDIO_DATA_DIR)
        if not os.path.exists(cls._AUDIO_LIST_DIR):
            os.makedirs(cls._AUDIO_LIST_DIR)

        with open(cls._JUNK_FILE, "w") as f:
            f.write("this is some garbage\nShould have no impact.")

        with open(cls._AUDIO_LIST_PATHS_PATH, "w") as f_list_fnames, \
                open(cls._AUDIO_LIST_FNAMES_PATH, "w") as f_list_paths:
            lengths = torch.randint(int(.5e5), int(1.5e6), (cls._N_EXAMPLES,))
            for i in range(cls._N_EXAMPLES):
                # dividing gets the noise in [-1, 1]
                white_noise = torch.randn((cls._N_CHANNELS, lengths[i])) / 10
                f_path = cls._AUDIO_DATA_PATH_FMT.format(i)
                torchaudio.save(f_path, white_noise, cls._SAMPLE_RATE)
                f_name_short = cls._AUDIO_DATA_FMT.format(i)
                f_list_fnames.write(f_name_short + "\n")
                f_list_paths.write(f_path + "\n") 

def setUpClass(cls):
        if not os.path.exists(cls._AUDIO_DATA_DIR):
            os.makedirs(cls._AUDIO_DATA_DIR)
        if not os.path.exists(cls._AUDIO_LIST_DIR):
            os.makedirs(cls._AUDIO_LIST_DIR)

        with open(cls._JUNK_FILE, "w") as f:
            f.write("this is some garbage\nShould have no impact.")

        with open(cls._AUDIO_LIST_PATHS_PATH, "w") as f_list_fnames, \
                open(cls._AUDIO_LIST_FNAMES_PATH, "w") as f_list_paths:
            lengths = torch.randint(int(.5e5), int(1.5e6), (cls._N_EXAMPLES,))
            for i in range(cls._N_EXAMPLES):
                # dividing gets the noise in [-1, 1]
                white_noise = torch.randn((cls._N_CHANNELS, lengths[i])) / 10
                f_path = cls._AUDIO_DATA_PATH_FMT.format(i)
                torchaudio.save(f_path, white_noise, cls._SAMPLE_RATE)
                f_name_short = cls._AUDIO_DATA_FMT.format(i)
                f_list_fnames.write(f_name_short + "\n")
                f_list_paths.write(f_path + "\n") 

def generate_background_noise(speech_commands):
    """Split the background noise provided by the dataset in 1 second chunks.

    Parameters:
        speech_commands (torch.utils.data.Dataset): Speech Command dataset as defined by torchaudio.
    """
    background_noise = glob.glob(
        os.path.join(speech_commands._path, "_background_noise_", "*.wav")
    )
    os.makedirs(os.path.join(speech_commands._path, "background"), exist_ok=True)

    for file in background_noise:
        waveform, sample_rate = torchaudio.load(file)
        background_waveforms = torch.split(waveform, sample_rate, dim=1)[:-1]

        for idx, background_waveform in enumerate(background_waveforms):
            torchaudio.save(
                os.path.join(
                    speech_commands._path,
                    "background",
                    f"{hash(waveform)}_nohash_{idx}.wav",
                ),
                background_waveform,
                sample_rate=sample_rate,
            ) 

def setUpClass(cls):
        if not os.path.exists(cls._AUDIO_DATA_DIR):
            os.makedirs(cls._AUDIO_DATA_DIR)
        if not os.path.exists(cls._AUDIO_LIST_DIR):
            os.makedirs(cls._AUDIO_LIST_DIR)

        with open(cls._JUNK_FILE, "w") as f:
            f.write("this is some garbage\nShould have no impact.")

        with open(cls._AUDIO_LIST_PATHS_PATH, "w") as f_list_fnames, \
                open(cls._AUDIO_LIST_FNAMES_PATH, "w") as f_list_paths:
            lengths = torch.randint(int(.5e5), int(1.5e6), (cls._N_EXAMPLES,))
            for i in range(cls._N_EXAMPLES):
                # dividing gets the noise in [-1, 1]
                white_noise = torch.randn((cls._N_CHANNELS, lengths[i])) / 10
                f_path = cls._AUDIO_DATA_PATH_FMT.format(i)
                torchaudio.save(f_path, white_noise, cls._SAMPLE_RATE)
                f_name_short = cls._AUDIO_DATA_FMT.format(i)
                f_list_fnames.write(f_name_short + "\n")
                f_list_paths.write(f_path + "\n") 

def _test_1_save(self, test_filepath, normalization):
        # load signal
        x, sr = torchaudio.load(test_filepath, normalization=normalization)

        # 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
        x = x[0, :]  # get mono signal
        x.squeeze_()  # remove channel dim
        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):
            x.unsqueeze_(1)  # L x C not C x L
            torchaudio.save(new_filepath, x, sr)

        with self.assertRaises(ValueError):
            x.squeeze_()
            x.unsqueeze_(1)
            x.unsqueeze_(0)  # 1 x L x 1
            torchaudio.save(new_filepath, x, sr)

        # don't save to folders that don't exist
        with self.assertRaises(OSError):
            new_filepath = os.path.join(self.test_dirpath, "no-path",
                                        "test.wav")
            torchaudio.save(new_filepath, x, sr) 

def _test_1_save_sine(self):

        # save created file
        sinewave_filepath = os.path.join(self.test_dirpath, "assets",
                                         "sinewave.wav")
        sr = 16000
        freq = 440
        volume = 0.3

        y = (torch.cos(
            2 * math.pi * torch.arange(0, 4 * sr).float() * freq / sr))
        y.unsqueeze_(0)
        # y is between -1 and 1, so must scale
        y = (y * volume * (2**31)).long()
        torchaudio.save(sinewave_filepath, y, sr)
        self.assertTrue(os.path.isfile(sinewave_filepath))

        # test precision
        new_precision = 32
        new_filepath = os.path.join(self.test_dirpath, "test.wav")
        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)
        os.unlink(new_filepath) 

def _test_3_load_and_save_is_identity(self):
        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) 

def _test_3_load_and_save_is_identity_across_backend(self, backend1, backend2):
        torchaudio.set_audio_backend(backend1)
        input_path = os.path.join(self.test_dirpath, 'assets', 'sinewave.wav')
        tensor1, sample_rate1 = torchaudio.load(input_path)

        output_path = os.path.join(self.test_dirpath, 'test.wav')
        torchaudio.save(output_path, tensor1, sample_rate1)

        torchaudio.set_audio_backend(backend2)
        tensor2, sample_rate2 = torchaudio.load(output_path)

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

def _test_4_load_partial(self):
        num_frames = 101
        offset = 201
        # 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)
        l1_error = x_sine_full[:, offset:(num_frames + offset)].sub(x_sine_part).abs().sum().item()
        # test for the correct number of samples and that the correct portion was loaded
        self.assertEqual(x_sine_part.size(1), num_frames)
        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)
        l1_error = x_2ch_sine_load.sub(x_2ch_sine[:, offset:(offset + num_frames)]).abs().sum().item()
        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)
        l1_error = x_2ch_full[:, offset:(offset + num_frames)].sub(x_2ch_part).abs().sum().item()
        self.assertEqual(x_2ch_part.size(1), num_frames)
        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)
        self.assertEqual(x_ns.size(1), x_sine_full.size(1) - offset_ns)

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

def py_save_func(
        filepath: str,
        tensor: torch.Tensor,
        sample_rate: int,
        channels_first: bool = True,
        compression: Optional[float] = None,
):
    torchaudio.save(filepath, tensor, sample_rate, channels_first, compression) 

def test_info_wav(self, dtype, sample_rate, num_channels):
        """`sox_io_backend.info` is torchscript-able and returns the same result"""
        audio_path = self.get_temp_path(f'{dtype}_{sample_rate}_{num_channels}.wav')
        data = get_wav_data(dtype, num_channels, normalize=False, num_frames=1 * sample_rate)
        save_wav(audio_path, data, sample_rate)

        script_path = self.get_temp_path('info_func.zip')
        torch.jit.script(py_info_func).save(script_path)
        ts_info_func = torch.jit.load(script_path)

        py_info = py_info_func(audio_path)
        ts_info = ts_info_func(audio_path)

        assert py_info.get_sample_rate() == ts_info.get_sample_rate()
        assert py_info.get_num_frames() == ts_info.get_num_frames()
        assert py_info.get_num_channels() == ts_info.get_num_channels()