| |
| |
|
|
| |
| |
|
|
| import math |
| import os |
| import random |
| import torch |
| import torch.utils.data |
| import numpy as np |
| from librosa.util import normalize |
| from scipy.io.wavfile import read |
| from librosa.filters import mel as librosa_mel_fn |
| import pathlib |
| from tqdm import tqdm |
|
|
| MAX_WAV_VALUE = 32767.0 |
|
|
|
|
| def load_wav(full_path, sr_target): |
| sampling_rate, data = read(full_path) |
| if sampling_rate != sr_target: |
| raise RuntimeError( |
| f"Sampling rate of the file {full_path} is {sampling_rate} Hz, but the model requires {sr_target} Hz" |
| ) |
| return data, sampling_rate |
|
|
|
|
| def dynamic_range_compression(x, C=1, clip_val=1e-5): |
| return np.log(np.clip(x, a_min=clip_val, a_max=None) * C) |
|
|
|
|
| def dynamic_range_decompression(x, C=1): |
| return np.exp(x) / C |
|
|
|
|
| def dynamic_range_compression_torch(x, C=1, clip_val=1e-5): |
| return torch.log(torch.clamp(x, min=clip_val) * C) |
|
|
|
|
| def dynamic_range_decompression_torch(x, C=1): |
| return torch.exp(x) / C |
|
|
|
|
| def spectral_normalize_torch(magnitudes): |
| return dynamic_range_compression_torch(magnitudes) |
|
|
|
|
| def spectral_de_normalize_torch(magnitudes): |
| return dynamic_range_decompression_torch(magnitudes) |
|
|
|
|
| mel_basis_cache = {} |
| hann_window_cache = {} |
|
|
|
|
| def mel_spectrogram( |
| y: torch.Tensor, |
| n_fft: int, |
| num_mels: int, |
| sampling_rate: int, |
| hop_size: int, |
| win_size: int, |
| fmin: int, |
| fmax: int = None, |
| center: bool = False, |
| ) -> torch.Tensor: |
| """ |
| Calculate the mel spectrogram of an input signal. |
| This function uses slaney norm for the librosa mel filterbank (using librosa.filters.mel) and uses Hann window for STFT (using torch.stft). |
| |
| Args: |
| y (torch.Tensor): Input signal. |
| n_fft (int): FFT size. |
| num_mels (int): Number of mel bins. |
| sampling_rate (int): Sampling rate of the input signal. |
| hop_size (int): Hop size for STFT. |
| win_size (int): Window size for STFT. |
| fmin (int): Minimum frequency for mel filterbank. |
| fmax (int): Maximum frequency for mel filterbank. If None, defaults to half the sampling rate (fmax = sr / 2.0) inside librosa_mel_fn |
| center (bool): Whether to pad the input to center the frames. Default is False. |
| |
| Returns: |
| torch.Tensor: Mel spectrogram. |
| """ |
| if torch.min(y) < -1.0: |
| print(f"[WARNING] Min value of input waveform signal is {torch.min(y)}") |
| if torch.max(y) > 1.0: |
| print(f"[WARNING] Max value of input waveform signal is {torch.max(y)}") |
|
|
| device = y.device |
| key = f"{n_fft}_{num_mels}_{sampling_rate}_{hop_size}_{win_size}_{fmin}_{fmax}_{device}" |
|
|
| if key not in mel_basis_cache: |
| mel = librosa_mel_fn( |
| sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax |
| ) |
| mel_basis_cache[key] = torch.from_numpy(mel).float().to(device) |
| hann_window_cache[key] = torch.hann_window(win_size).to(device) |
|
|
| mel_basis = mel_basis_cache[key] |
| hann_window = hann_window_cache[key] |
|
|
| padding = (n_fft - hop_size) // 2 |
| y = torch.nn.functional.pad( |
| y.unsqueeze(1), (padding, padding), mode="reflect" |
| ).squeeze(1) |
|
|
| spec = torch.stft( |
| y, |
| n_fft, |
| hop_length=hop_size, |
| win_length=win_size, |
| window=hann_window, |
| center=center, |
| pad_mode="reflect", |
| normalized=False, |
| onesided=True, |
| return_complex=True, |
| ) |
| spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9) |
|
|
| mel_spec = torch.matmul(mel_basis, spec) |
| mel_spec = spectral_normalize_torch(mel_spec) |
|
|
| return mel_spec |
|
|
|
|
| def get_mel_spectrogram(wav, h): |
| """ |
| Generate mel spectrogram from a waveform using given hyperparameters. |
| |
| Args: |
| wav (torch.Tensor): Input waveform. |
| h: Hyperparameters object with attributes n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax. |
| |
| Returns: |
| torch.Tensor: Mel spectrogram. |
| """ |
| return mel_spectrogram( |
| wav, |
| h.n_fft, |
| h.num_mels, |
| h.sampling_rate, |
| h.hop_size, |
| h.win_size, |
| h.fmin, |
| h.fmax, |
| ) |
|
|
|
|
| def get_dataset_filelist(a): |
| training_files = [] |
| validation_files = [] |
| list_unseen_validation_files = [] |
|
|
| with open(a.input_training_file, "r", encoding="utf-8") as fi: |
| training_files = [ |
| os.path.join(a.input_wavs_dir, x.split("|")[0] + ".wav") |
| for x in fi.read().split("\n") |
| if len(x) > 0 |
| ] |
| print(f"first training file: {training_files[0]}") |
|
|
| with open(a.input_validation_file, "r", encoding="utf-8") as fi: |
| validation_files = [ |
| os.path.join(a.input_wavs_dir, x.split("|")[0] + ".wav") |
| for x in fi.read().split("\n") |
| if len(x) > 0 |
| ] |
| print(f"first validation file: {validation_files[0]}") |
|
|
| for i in range(len(a.list_input_unseen_validation_file)): |
| with open(a.list_input_unseen_validation_file[i], "r", encoding="utf-8") as fi: |
| unseen_validation_files = [ |
| os.path.join(a.list_input_unseen_wavs_dir[i], x.split("|")[0] + ".wav") |
| for x in fi.read().split("\n") |
| if len(x) > 0 |
| ] |
| print( |
| f"first unseen {i}th validation fileset: {unseen_validation_files[0]}" |
| ) |
| list_unseen_validation_files.append(unseen_validation_files) |
|
|
| return training_files, validation_files, list_unseen_validation_files |
|
|
|
|
| class MelDataset(torch.utils.data.Dataset): |
| def __init__( |
| self, |
| training_files, |
| hparams, |
| segment_size, |
| n_fft, |
| num_mels, |
| hop_size, |
| win_size, |
| sampling_rate, |
| fmin, |
| fmax, |
| split=True, |
| shuffle=True, |
| n_cache_reuse=1, |
| device=None, |
| fmax_loss=None, |
| fine_tuning=False, |
| base_mels_path=None, |
| is_seen=True, |
| ): |
| self.audio_files = training_files |
| random.seed(1234) |
| if shuffle: |
| random.shuffle(self.audio_files) |
| self.hparams = hparams |
| self.is_seen = is_seen |
| if self.is_seen: |
| self.name = pathlib.Path(self.audio_files[0]).parts[0] |
| else: |
| self.name = "-".join(pathlib.Path(self.audio_files[0]).parts[:2]).strip("/") |
|
|
| self.segment_size = segment_size |
| self.sampling_rate = sampling_rate |
| self.split = split |
| self.n_fft = n_fft |
| self.num_mels = num_mels |
| self.hop_size = hop_size |
| self.win_size = win_size |
| self.fmin = fmin |
| self.fmax = fmax |
| self.fmax_loss = fmax_loss |
| self.cached_wav = None |
| self.n_cache_reuse = n_cache_reuse |
| self._cache_ref_count = 0 |
| self.device = device |
| self.fine_tuning = fine_tuning |
| self.base_mels_path = base_mels_path |
|
|
| print("[INFO] checking dataset integrity...") |
| for i in tqdm(range(len(self.audio_files))): |
| assert os.path.exists( |
| self.audio_files[i] |
| ), f"{self.audio_files[i]} not found" |
|
|
| def __getitem__(self, index): |
| filename = self.audio_files[index] |
| if self._cache_ref_count == 0: |
| audio, sampling_rate = load_wav(filename, self.sampling_rate) |
| audio = audio / MAX_WAV_VALUE |
| if not self.fine_tuning: |
| audio = normalize(audio) * 0.95 |
| self.cached_wav = audio |
| if sampling_rate != self.sampling_rate: |
| raise ValueError( |
| f"{sampling_rate} SR doesn't match target {self.sampling_rate} SR" |
| ) |
| self._cache_ref_count = self.n_cache_reuse |
| else: |
| audio = self.cached_wav |
| self._cache_ref_count -= 1 |
|
|
| audio = torch.FloatTensor(audio) |
| audio = audio.unsqueeze(0) |
|
|
| if not self.fine_tuning: |
| if self.split: |
| if audio.size(1) >= self.segment_size: |
| max_audio_start = audio.size(1) - self.segment_size |
| audio_start = random.randint(0, max_audio_start) |
| audio = audio[:, audio_start : audio_start + self.segment_size] |
| else: |
| audio = torch.nn.functional.pad( |
| audio, (0, self.segment_size - audio.size(1)), "constant" |
| ) |
|
|
| mel = mel_spectrogram( |
| audio, |
| self.n_fft, |
| self.num_mels, |
| self.sampling_rate, |
| self.hop_size, |
| self.win_size, |
| self.fmin, |
| self.fmax, |
| center=False, |
| ) |
| else: |
| |
| if (audio.size(1) % self.hop_size) != 0: |
| audio = audio[:, : -(audio.size(1) % self.hop_size)] |
| mel = mel_spectrogram( |
| audio, |
| self.n_fft, |
| self.num_mels, |
| self.sampling_rate, |
| self.hop_size, |
| self.win_size, |
| self.fmin, |
| self.fmax, |
| center=False, |
| ) |
| assert ( |
| audio.shape[1] == mel.shape[2] * self.hop_size |
| ), f"audio shape {audio.shape} mel shape {mel.shape}" |
|
|
| else: |
| mel = np.load( |
| os.path.join( |
| self.base_mels_path, |
| os.path.splitext(os.path.split(filename)[-1])[0] + ".npy", |
| ) |
| ) |
| mel = torch.from_numpy(mel) |
|
|
| if len(mel.shape) < 3: |
| mel = mel.unsqueeze(0) |
|
|
| if self.split: |
| frames_per_seg = math.ceil(self.segment_size / self.hop_size) |
|
|
| if audio.size(1) >= self.segment_size: |
| mel_start = random.randint(0, mel.size(2) - frames_per_seg - 1) |
| mel = mel[:, :, mel_start : mel_start + frames_per_seg] |
| audio = audio[ |
| :, |
| mel_start |
| * self.hop_size : (mel_start + frames_per_seg) |
| * self.hop_size, |
| ] |
| else: |
| mel = torch.nn.functional.pad( |
| mel, (0, frames_per_seg - mel.size(2)), "constant" |
| ) |
| audio = torch.nn.functional.pad( |
| audio, (0, self.segment_size - audio.size(1)), "constant" |
| ) |
|
|
| mel_loss = mel_spectrogram( |
| audio, |
| self.n_fft, |
| self.num_mels, |
| self.sampling_rate, |
| self.hop_size, |
| self.win_size, |
| self.fmin, |
| self.fmax_loss, |
| center=False, |
| ) |
|
|
| return (mel.squeeze(), audio.squeeze(0), filename, mel_loss.squeeze()) |
|
|
| def __len__(self): |
| return len(self.audio_files) |
|
|
