import math, os, re, sys from pathlib import Path import numpy as np import pandas as pd import Levenshtein from multiprocessing import Pool from scipy.io import wavfile import tensorflow as tf from tensorflow.keras.utils import Sequence, OrderedEnqueuer from tensorflow.keras import layers from tensorflow.keras.preprocessing.sequence import pad_sequences sys.path.append(os.path.dirname(__file__)) from g2p.g2p_en.g2p import G2p import warnings warnings.filterwarnings("ignore", category=np.VisibleDeprecationWarning) np.warnings.filterwarnings('ignore', category=np.VisibleDeprecationWarning) class LibriPhraseDataloader(Sequence): def __init__(self, batch_size, fs = 16000, wav_dir='/homw/DB/LibriPhrase/wav_dir', noise_dir='/homw/DB/noise', csv_dir='/homw/DB/LibriPhrase/data', train_csv = ['train_100h', 'train_360h'], test_csv = ['train_500h',], types='both', # easy, hard features='g2p_embed', # phoneme, g2p_embed, both ... train=True, shuffle=True, pkl=None, edit_dist=False, ): phonemes = ["", ] + ['AA0', 'AA1', 'AA2', 'AE0', 'AE1', 'AE2', 'AH0', 'AH1', 'AH2', 'AO0', 'AO1', 'AO2', 'AW0', 'AW1', 'AW2', 'AY0', 'AY1', 'AY2', 'B', 'CH', 'D', 'DH', 'EH0', 'EH1', 'EH2', 'ER0', 'ER1', 'ER2', 'EY0', 'EY1', 'EY2', 'F', 'G', 'HH', 'IH0', 'IH1', 'IH2', 'IY0', 'IY1', 'IY2', 'JH', 'K', 'L', 'M', 'N', 'NG', 'OW0', 'OW1', 'OW2', 'OY0', 'OY1', 'OY2', 'P', 'R', 'S', 'SH', 'T', 'TH', 'UH0', 'UH1', 'UH2', 'UW', 'UW0', 'UW1', 'UW2', 'V', 'W', 'Y', 'Z', 'ZH', ' '] self.p2idx = {p: idx for idx, p in enumerate(phonemes)} self.idx2p = {idx: p for idx, p in enumerate(phonemes)} self.batch_size = batch_size self.fs = fs self.wav_dir = wav_dir self.csv_dir = csv_dir self.noise_dir = noise_dir self.train_csv = train_csv self.test_csv = test_csv self.types = types self.features = features self.train = train self.shuffle = shuffle self.pkl = pkl self.edit_dist = edit_dist self.nPhoneme = len(phonemes) self.g2p = G2p() self.__prep__() self.on_epoch_end() def __prep__(self): if self.train: print(">> Preparing noise DB") noise_list = [str(x) for x in Path(self.noise_dir).rglob('*.wav')] self.noise = np.array([]) for noise in noise_list: fs, data = wavfile.read(noise) assert fs == self.fs, ">> Error : Un-match sampling freq.\n{} -> {}".format(noise, fs) data = data.astype(np.float32) / 32768.0 data = (data / np.max(data)) * 0.5 self.noise = np.append(self.noise, data) self.data = pd.DataFrame(columns=['wav_label', 'wav', 'text', 'duration', 'label', 'type']) if (self.pkl is not None) and (os.path.isfile(self.pkl)): print(">> Load dataset from {}".format(self.pkl)) self.data = pd.read_pickle(self.pkl) else: for db in self.train_csv if self.train else self.test_csv: csv_list = [str(x) for x in Path(self.csv_dir).rglob('*' + db + '*word*')] for n_word in csv_list: print(">> processing : {} ".format(n_word)) df = pd.read_csv(n_word) # Split train dataset to match & unmatch case anc_pos = df[['anchor_text', 'anchor', 'anchor_text', 'anchor_dur']] anc_neg = df[['anchor_text', 'anchor', 'comparison_text', 'anchor_dur', 'target', 'type']] com_pos = df[['comparison_text', 'comparison', 'comparison_text', 'comparison_dur']] com_neg = df[['comparison_text', 'comparison', 'anchor_text', 'comparison_dur', 'target', 'type']] anc_pos.columns = ['wav_label', 'anchor', 'anchor_text', 'anchor_dur'] com_pos.columns = ['wav_label', 'comparison', 'comparison_text', 'comparison_dur'] anc_pos['label'] = 1 anc_pos['type'] = df['type'] com_pos['label'] = 1 com_pos['type'] = df['type'] # Concat self.data = self.data.append(anc_pos.rename(columns={y: x for x, y in zip(self.data.columns, anc_pos.columns)}), ignore_index=True) self.data = self.data.append(anc_neg.rename(columns={y: x for x, y in zip(self.data.columns, anc_neg.columns)}), ignore_index=True) self.data = self.data.append(com_pos.rename(columns={y: x for x, y in zip(self.data.columns, com_pos.columns)}), ignore_index=True) self.data = self.data.append(com_neg.rename(columns={y: x for x, y in zip(self.data.columns, com_neg.columns)}), ignore_index=True) # Append wav directory path self.data['wav'] = self.data['wav'].apply(lambda x: os.path.join(self.wav_dir, x)) # g2p & p2idx by g2p_en package print(">> Convert word to phoneme") self.data['phoneme'] = self.data['text'].apply(lambda x: self.g2p(re.sub(r"[^a-zA-Z0-9]+", ' ', x))) print(">> Convert speech word to phoneme") self.data['wav_phoneme'] = self.data['wav_label'].apply(lambda x: self.g2p(re.sub(r"[^a-zA-Z0-9]+", ' ', x))) print(">> Convert phoneme to index") self.data['pIndex'] = self.data['phoneme'].apply(lambda x: [self.p2idx[t] for t in x]) print(">> Convert speech phoneme to index") self.data['wav_pIndex'] = self.data['wav_phoneme'].apply(lambda x: [self.p2idx[t] for t in x]) print(">> Compute phoneme embedding") self.data['g2p_embed'] = self.data['text'].apply(lambda x: self.g2p.embedding(x)) print(">> Calucate Edit distance ratio") self.data['dist'] = self.data.apply(lambda x: Levenshtein.ratio(re.sub(r"[^a-zA-Z0-9]+", ' ', x['wav_label']), re.sub(r"[^a-zA-Z0-9]+", ' ', x['text'])), axis=1) if (self.pkl is not None) and (not os.path.isfile(self.pkl)): self.data.to_pickle(self.pkl) # Masking dataset type if self.types == 'both': pass elif self.types == 'easy': self.data = self.data.loc[self.data['type'] == 'diffspk_easyneg'] elif self.types == 'hard': self.data = self.data.loc[self.data['type'] == 'diffspk_hardneg'] # Get longest data self.data = self.data.sort_values(by='duration').reset_index(drop=True) self.wav_list = self.data['wav'].values self.idx_list = self.data['pIndex'].values self.sIdx_list = self.data['wav_pIndex'].values self.emb_list = self.data['g2p_embed'].values self.lab_list = self.data['label'].values if self.edit_dist: self.dist_list = self.data['dist'].values # Set dataloader params. self.len = len(self.data) self.maxlen_t = int((int(self.data['text'].apply(lambda x: len(x)).max() / 10) + 1) * 10) self.maxlen_a = int((int(self.data['duration'].values[-1] / 0.5) + 1 ) * self.fs / 2) self.maxlen_l = int((int(self.data['wav_label'].apply(lambda x: len(x)).max() / 10) + 1) * 10) def __len__(self): # return total batch-wise length return math.ceil(self.len / self.batch_size) def _load_wav(self, wav): return np.array(wavfile.read(wav)[1]).astype(np.float32) / 32768.0 def _mixing_snr(self, clean, snr=[5, 15]): def _cal_adjusted_rms(clean_rms, snr): a = float(snr) / 20 noise_rms = clean_rms / (10**a) return noise_rms def _cal_rms(amp): return np.sqrt(np.mean(np.square(amp), axis=-1)) start = np.random.randint(0, len(self.noise)-len(clean)) divided_noise = self.noise[start: start + len(clean)] clean_rms = _cal_rms(clean) noise_rms = _cal_rms(divided_noise) adj_noise_rms = _cal_adjusted_rms(clean_rms, np.random.randint(snr[0], snr[1])) adj_noise_amp = divided_noise * (adj_noise_rms / (noise_rms + 1e-7)) noisy = clean + adj_noise_amp if np.max(noisy) > 1: noisy = noisy / np.max(noisy) return noisy def __getitem__(self, idx): # chunking indices = self.indices[idx * self.batch_size : (idx + 1) * self.batch_size] # load inputs batch_x = [np.array(wavfile.read(self.wav_list[i])[1]).astype(np.float32) / 32768.0 for i in indices] if self.features == 'both': batch_p = [np.array(self.idx_list[i]).astype(np.int32) for i in indices] batch_e = [np.array(self.emb_list[i]).astype(np.float32) for i in indices] else: if self.features == 'phoneme': batch_y = [np.array(self.idx_list[i]).astype(np.int32) for i in indices] elif self.features == 'g2p_embed': batch_y = [np.array(self.emb_list[i]).astype(np.float32) for i in indices] # load outputs batch_z = [np.array([self.lab_list[i]]).astype(np.float32) for i in indices] batch_l = [np.array(self.sIdx_list[i]).astype(np.int32) for i in indices] batch_t = [np.array(self.idx_list[i]).astype(np.int32) for i in indices] if self.edit_dist: batch_d = [np.array([self.dist_list[i]]).astype(np.float32) for i in indices] # padding and masking pad_batch_x = pad_sequences(np.array(batch_x), maxlen=self.maxlen_a, value=0.0, padding='post', dtype=batch_x[0].dtype) if self.features == 'both': pad_batch_p = pad_sequences(np.array(batch_p), maxlen=self.maxlen_t, value=0.0, padding='post', dtype=batch_p[0].dtype) pad_batch_e = pad_sequences(np.array(batch_e), maxlen=self.maxlen_t, value=0.0, padding='post', dtype=batch_e[0].dtype) else: pad_batch_y = pad_sequences(np.array(batch_y), maxlen=self.maxlen_t, value=0.0, padding='post', dtype=batch_y[0].dtype) pad_batch_z = pad_sequences(np.array(batch_z), value=0.0, padding='post', dtype=batch_z[0].dtype) pad_batch_l = pad_sequences(np.array(batch_l), maxlen=self.maxlen_l, value=0.0, padding='post', dtype=batch_l[0].dtype) pad_batch_t = pad_sequences(np.array(batch_t), maxlen=self.maxlen_t, value=0.0, padding='post', dtype=batch_t[0].dtype) if self.edit_dist: pad_batch_d = pad_sequences(np.array(batch_d), value=0.0, padding='post', dtype=batch_d[0].dtype) # Noisy option if self.train: batch_x_noisy = [self._mixing_snr(x) for x in batch_x] pad_batch_x_noisy = pad_sequences(np.array(batch_x_noisy), maxlen=self.maxlen_a, value=0.0, padding='post', dtype=batch_x_noisy[0].dtype) if self.train: if self.features == 'both': return pad_batch_x, pad_batch_x_noisy, pad_batch_p, pad_batch_e, pad_batch_z, pad_batch_l, pad_batch_t else: return pad_batch_x, pad_batch_x_noisy, pad_batch_y, pad_batch_z, pad_batch_l, pad_batch_t else: if self.features == 'both': if self.edit_dist: return pad_batch_x, pad_batch_p, pad_batch_e, pad_batch_z, pad_batch_d else: return pad_batch_x, pad_batch_p, pad_batch_e, pad_batch_z else: if self.edit_dist: return pad_batch_x, pad_batch_y, pad_batch_z, pad_batch_d else: return pad_batch_x, pad_batch_y, pad_batch_z def on_epoch_end(self): self.indices = np.arange(self.len) if self.shuffle == True: np.random.shuffle(self.indices) def convert_sequence_to_dataset(dataloader): def data_generator(): for i in range(dataloader.__len__()): if dataloader.train: if dataloader.features == 'both': pad_batch_x, pad_batch_x_noisy, pad_batch_p, pad_batch_e, pad_batch_z, pad_batch_l, pad_batch_t = dataloader[i] yield pad_batch_x, pad_batch_x_noisy, pad_batch_p, pad_batch_e, pad_batch_z, pad_batch_l, pad_batch_t else: pad_batch_x, pad_batch_x_noisy, pad_batch_y, pad_batch_z, pad_batch_l, pad_batch_t = dataloader[i] yield pad_batch_x, pad_batch_x_noisy, pad_batch_y, pad_batch_z, pad_batch_l, pad_batch_t else: if dataloader.features == 'both': if dataloader.edit_dist: pad_batch_x, pad_batch_p, pad_batch_e, pad_batch_z, pad_batch_d = dataloader[i] yield pad_batch_x, pad_batch_p, pad_batch_e, pad_batch_z, pad_batch_d else: pad_batch_x, pad_batch_p, pad_batch_e, pad_batch_z = dataloader[i] yield pad_batch_x, pad_batch_p, pad_batch_e, pad_batch_z else: if dataloader.edit_dist: pad_batch_x, pad_batch_y, pad_batch_z, pad_batch_d = dataloader[i] yield pad_batch_x, pad_batch_y, pad_batch_z, pad_batch_d else: pad_batch_x, pad_batch_y, pad_batch_z = dataloader[i] yield pad_batch_x, pad_batch_y, pad_batch_z if dataloader.train: if dataloader.features == 'both': data_dataset = tf.data.Dataset.from_generator(data_generator, output_signature=( tf.TensorSpec(shape=(None, dataloader.maxlen_a), dtype=tf.float32), tf.TensorSpec(shape=(None, dataloader.maxlen_a), dtype=tf.float32), tf.TensorSpec(shape=(None, dataloader.maxlen_t), dtype=tf.int32), tf.TensorSpec(shape=(None, dataloader.maxlen_t, 256), dtype=tf.float32), tf.TensorSpec(shape=(None, 1), dtype=tf.float32), tf.TensorSpec(shape=(None, dataloader.maxlen_l), dtype=tf.int32), tf.TensorSpec(shape=(None, dataloader.maxlen_t), dtype=tf.int32),) ) else: data_dataset = tf.data.Dataset.from_generator(data_generator, output_signature=( tf.TensorSpec(shape=(None, dataloader.maxlen_a), dtype=tf.float32), tf.TensorSpec(shape=(None, dataloader.maxlen_a), dtype=tf.float32), tf.TensorSpec(shape=(None, dataloader.maxlen_t) if dataloader.features == 'phoneme' else (None, dataloader.maxlen_t, 256), dtype=tf.int32 if dataloader.features == 'phoneme' else tf.float32), tf.TensorSpec(shape=(None, 1), dtype=tf.float32), tf.TensorSpec(shape=(None, dataloader.maxlen_l), dtype=tf.int32), tf.TensorSpec(shape=(None, dataloader.maxlen_t), dtype=tf.int32),) ) else: if dataloader.features == 'both': if dataloader.edit_dist: data_dataset = tf.data.Dataset.from_generator(data_generator, output_signature=( tf.TensorSpec(shape=(None, dataloader.maxlen_a), dtype=tf.float32), tf.TensorSpec(shape=(None, dataloader.maxlen_t), dtype=tf.int32), tf.TensorSpec(shape=(None, dataloader.maxlen_t, 256), dtype=tf.float32), tf.TensorSpec(shape=(None, 1), dtype=tf.float32), tf.TensorSpec(shape=(None, 1), dtype=tf.float32),) ) else: data_dataset = tf.data.Dataset.from_generator(data_generator, output_signature=( tf.TensorSpec(shape=(None, dataloader.maxlen_a), dtype=tf.float32), tf.TensorSpec(shape=(None, dataloader.maxlen_t), dtype=tf.int32), tf.TensorSpec(shape=(None, dataloader.maxlen_t, 256), dtype=tf.float32), tf.TensorSpec(shape=(None, 1), dtype=tf.float32),) ) else: if dataloader.edit_dist: data_dataset = tf.data.Dataset.from_generator(data_generator, output_signature=( tf.TensorSpec(shape=(None, dataloader.maxlen_a), dtype=tf.float32), tf.TensorSpec(shape=(None, dataloader.maxlen_t) if dataloader.features == 'phoneme' else (None, dataloader.maxlen_t, 256), dtype=tf.int32 if dataloader.features == 'phoneme' else tf.float32), tf.TensorSpec(shape=(None, 1), dtype=tf.float32), tf.TensorSpec(shape=(None, 1), dtype=tf.float32),) ) else: data_dataset = tf.data.Dataset.from_generator(data_generator, output_signature=( tf.TensorSpec(shape=(None, dataloader.maxlen_a), dtype=tf.float32), tf.TensorSpec(shape=(None, dataloader.maxlen_t) if dataloader.features == 'phoneme' else (None, dataloader.maxlen_t, 256), dtype=tf.int32 if dataloader.features == 'phoneme' else tf.float32), tf.TensorSpec(shape=(None, 1), dtype=tf.float32),) ) # data_dataset = data_dataset.cache() data_dataset = data_dataset.prefetch(1) return data_dataset if __name__ == '__main__': GLOBAL_BATCH_SIZE = 2048 train_dataset = LibriPhraseDataloader(batch_size=GLOBAL_BATCH_SIZE, train=True, types='both', shuffle=True, pkl='/home/DB/LibriPhrase/data/train_both.pkl', features='g2p_embed') test_dataset = LibriPhraseDataloader(batch_size=GLOBAL_BATCH_SIZE, train=False, edit_dist=True, types='both', shuffle=False, pkl='/home/DB/LibriPhrase/data/test_both.pkl', features='g2p_embed')