AuxiliaryASR/meldataset.py

# coding: utf-8

import os
import os.path as osp
import time
import random
import numpy as np
import random

import soundfile
import soundfile as sf

import torch
from torch import nn
import torch.nn.functional as F
import torchaudio
from torch.utils.data import DataLoader
from datasets import load_dataset

from g2p_en import G2p
from tqdm import tqdm

import logging

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)
from text_utils import TextCleaner

np.random.seed(1)
random.seed(1)
DEFAULT_DICT_PATH = osp.join(osp.dirname(__file__), 'word_index_dict.csv')
SPECT_PARAMS = {
    "n_fft": 2048,
    "win_length": 2048,
    "hop_length": 512
}
MEL_PARAMS = {
    "n_mels": 80,
    "n_fft": 2048,
    "win_length": 2048,
    "hop_length": 512
}


class MelDataset(torch.utils.data.Dataset):
    def __init__(self,
                 data_list,
                 dict_path=DEFAULT_DICT_PATH,
                 sr=24000,
                 spect_params=SPECT_PARAMS,
                 mel_params=MEL_PARAMS,
                 ):

        _data_list = [l[:-1].split('|') for l in data_list]
        self.data_list = [data if len(data) == 3 else (*data, 0) for data in _data_list]
        self.text_cleaner = TextCleaner(dict_path)
        self.sr = sr

        self.to_melspec = torchaudio.transforms.MelSpectrogram(**mel_params, sample_rate=sr, pad_mode="reflect")
        self.mean, self.std = -4, 4

    def __len__(self):
        return len(self.data_list)

    def __getitem__(self, idx):
        data = self.data_list[idx]
        wave, text_tensor, speaker_id = self._load_tensor(data)
        wave_tensor = torch.from_numpy(wave).float()
        mel_tensor = self.to_melspec(wave_tensor)
        if (text_tensor.size(0) + 1) >= (mel_tensor.size(1) // 3):
            # if len(mel_tensor.size()) != 3:
            tenst = mel_tensor.unsqueeze(0)
            # print(tenst.size())
            # print((text_tensor.size(0)+1)*3)
            # print(data)
            # else:
            #     tenst = mel_tensor
            mel_tensor = F.interpolate(
                tenst, size=(text_tensor.size(0) + 1) * 3, align_corners=False,
                mode='linear').squeeze(0)

        acoustic_feature = (torch.log(1e-5 + mel_tensor) - self.mean) / self.std

        length_feature = acoustic_feature.size(1)
        acoustic_feature = acoustic_feature[:, :(length_feature - length_feature % 2)]

        return wave_tensor, acoustic_feature, text_tensor, data[0]

    def _load_tensor(self, data):
        wave_path, text, speaker_id = data
        wave_path = os.path.normpath(wave_path)
        speaker_id = int(speaker_id)
        wave, sr = sf.read(wave_path)

        # Convert to mono if needed
        if np.ndim(wave) > 1:
            wave = np.mean(wave, axis=1)

        # phonemize the text
        ps = text
        text = self.text_cleaner(ps)
        blank_index = self.text_cleaner.word_index_dictionary[" "]
        text.insert(0, blank_index)  # add a blank at the beginning (silence)
        text.append(blank_index)  # add a blank at the end (silence)

        text = torch.LongTensor(text)

        return wave, text, speaker_id


class HFMelDataset(torch.utils.data.Dataset):
    def __init__(self,
                 data_name,
                 split="train",
                 audio_column="audio",
                 text_column="text",
                 speaker_column="text",
                 dict_path=DEFAULT_DICT_PATH,
                 sr=24000,
                 spect_params=SPECT_PARAMS,
                 mel_params=MEL_PARAMS,
                 ):
        self.dataset = load_dataset(data_name, split=split)
        self.audio_column = audio_column
        self.text_column = text_column
        self.speaker_column = speaker_column

        self.text_cleaner = TextCleaner(dict_path)
        self.sr = sr

        self.to_melspec = torchaudio.transforms.MelSpectrogram(**mel_params, sample_rate=sr, pad_mode="reflect")
        self.mean, self.std = -4, 4

    def __len__(self):
        return len(self.data_list)

    def __getitem__(self, idx):
        data = self.dataset[idx]
        wave, _, speaker_id = (data[self.audio_column], data[self.text_column], data[self.speaker_column])

        ps = data[self.text_column]
        text = self.text_cleaner(ps)
        blank_index = self.text_cleaner.word_index_dictionary[" "]
        text.insert(0, blank_index)  # add a blank at the beginning (silence)
        text.append(blank_index)  # add a blank at the end (silence)

        text_tensor = torch.LongTensor(text)

        mel_tensor = self.to_melspec(wave)
        if (text_tensor.size(0) + 1) >= (mel_tensor.size(1) // 3):
            # if len(mel_tensor.size()) != 3:
            tenst = mel_tensor.unsqueeze(0)
            # print(tenst.size())
            # print((text_tensor.size(0)+1)*3)
            # print(data)
            # else:
            #     tenst = mel_tensor
            mel_tensor = F.interpolate(
                tenst, size=(text_tensor.size(0) + 1) * 3, align_corners=False,
                mode='linear').squeeze(0)

        acoustic_feature = (torch.log(1e-5 + mel_tensor) - self.mean) / self.std

        length_feature = acoustic_feature.size(1)
        acoustic_feature = acoustic_feature[:, :(length_feature - length_feature % 2)]

        return wave, acoustic_feature, text_tensor, data[0]


class Collater(object):
    """
    Args:
      return_wave (bool): if true, will return the wave data along with spectrogram. 
    """

    def __init__(self, return_wave=False):
        self.text_pad_index = 0
        self.return_wave = return_wave

    def __call__(self, batch):
        batch_size = len(batch)

        # sort by mel length
        lengths = [b[1].shape[1] for b in batch]
        batch_indexes = np.argsort(lengths)[::-1]
        batch = [batch[bid] for bid in batch_indexes]

        nmels = batch[0][1].size(0)
        max_mel_length = max([b[1].shape[1] for b in batch])
        max_text_length = max([b[2].shape[0] for b in batch])

        mels = torch.zeros((batch_size, nmels, max_mel_length)).float()
        texts = torch.zeros((batch_size, max_text_length)).long()
        input_lengths = torch.zeros(batch_size).long()
        output_lengths = torch.zeros(batch_size).long()
        paths = ['' for _ in range(batch_size)]
        for bid, (_, mel, text, path) in enumerate(batch):
            mel_size = mel.size(1)
            text_size = text.size(0)
            mels[bid, :, :mel_size] = mel
            texts[bid, :text_size] = text
            input_lengths[bid] = text_size
            output_lengths[bid] = mel_size
            paths[bid] = path
            assert (text_size < (mel_size // 2))

        if self.return_wave:
            waves = [b[0] for b in batch]
            return texts, input_lengths, mels, output_lengths, paths, waves

        return texts, input_lengths, mels, output_lengths


def build_dataloader(path_list,
                     validation=False,
                     batch_size=4,
                     device='cpu',
                     collate_config={},
                     dataset_config={}):
    dataset = MelDataset(path_list, **dataset_config)
    collate_fn = Collater(**collate_config)
    data_loader = DataLoader(dataset,
                             batch_size=batch_size,
                             shuffle=(not validation),
                             drop_last=(not validation),
                             collate_fn=collate_fn,
                             pin_memory=(device != 'cpu'))

    return data_loader


def build_dataloaderHF(name,
                       split="train",
                       audio_column="audio",
                       text_column="text",
                       speaker_column="text",
                       validation=False,
                       batch_size=4,
                       device='cpu',
                       collate_config={},
                       dataset_config={}):

    dataset = HFMelDataset(name, split=split,
                           audio_column=audio_column,
                           text_column=text_column,
                           speaker_column=speaker_column,
                           **dataset_config)
    collate_fn = Collater(**collate_config)
    data_loader = DataLoader(dataset,
                             batch_size=batch_size,
                             shuffle=(not validation),
                             drop_last=(not validation),
                             collate_fn=collate_fn,
                             pin_memory=(device != 'cpu'))

    return data_loader