import json import pathlib import pickle import random import shutil import warnings from copy import deepcopy import numpy as np import torch from tqdm import tqdm from utils.hparams import hparams from utils.indexed_datasets import IndexedDatasetBuilder from utils.multiprocess_utils import chunked_multiprocess_run from utils.phoneme_utils import load_phoneme_dictionary from utils.plot import distribution_to_figure class BinarizationError(Exception): pass class BaseBinarizer: """ Base class for data processing. 1. *process* and *process_data_split*: process entire data, generate the train-test split (support parallel processing); 2. *process_item*: process singe piece of data; 3. *get_pitch*: infer the pitch using some algorithm; 4. *get_align*: get the alignment using 'mel2ph' format (see https://arxiv.org/abs/1905.09263). 5. phoneme encoder, voice encoder, etc. Subclasses should define: 1. *load_metadata*: how to read multiple datasets from files; 2. *train_item_names*, *valid_item_names*, *test_item_names*: how to split the dataset; 3. load_ph_set: the phoneme set. """ def __init__(self, datasets=None, data_attrs=None): if datasets is None: datasets = hparams['datasets'] self.datasets = datasets self.raw_data_dirs = [pathlib.Path(ds['raw_data_dir']) for ds in self.datasets] self.binary_data_dir = pathlib.Path(hparams['binary_data_dir']) self.data_attrs = [] if data_attrs is None else data_attrs self.binarization_args = hparams['binarization_args'] self.augmentation_args = hparams.get('augmentation_args', {}) self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') self.spk_map = {} self.spk_ids = None self.build_spk_map() self.lang_map = {} self.dictionaries = hparams['dictionaries'] self.build_lang_map() self.items = {} self.item_names: list = None self._train_item_names: list = None self._valid_item_names: list = None self.phoneme_dictionary = load_phoneme_dictionary() self.timestep = hparams['hop_size'] / hparams['audio_sample_rate'] def build_spk_map(self): spk_ids = [ds.get('spk_id') for ds in self.datasets] assigned_spk_ids = {spk_id for spk_id in spk_ids if spk_id is not None} idx = 0 for i in range(len(spk_ids)): if spk_ids[i] is not None: continue while idx in assigned_spk_ids: idx += 1 spk_ids[i] = idx assigned_spk_ids.add(idx) assert max(spk_ids) < hparams['num_spk'], \ f'Index in spk_id sequence {spk_ids} is out of range. All values should be smaller than num_spk.' for spk_id, dataset in zip(spk_ids, self.datasets): spk_name = dataset['speaker'] if spk_name in self.spk_map and self.spk_map[spk_name] != spk_id: raise ValueError(f'Invalid speaker ID assignment. Name \'{spk_name}\' is assigned ' f'with different speaker IDs: {self.spk_map[spk_name]} and {spk_id}.') self.spk_map[spk_name] = spk_id self.spk_ids = spk_ids print("| spk_map: ", self.spk_map) def build_lang_map(self): assert len(self.dictionaries.keys()) <= hparams['num_lang'], \ 'Number of languages must not be greater than num_lang!' for dataset in self.datasets: assert dataset['language'] in self.dictionaries, f'Unrecognized language name: {dataset["language"]}' for lang_id, lang_name in enumerate(sorted(self.dictionaries.keys()), start=1): self.lang_map[lang_name] = lang_id print("| lang_map: ", self.lang_map) def load_meta_data(self, raw_data_dir: pathlib.Path, ds_id, spk, lang) -> dict: raise NotImplementedError() def split_train_valid_set(self, prefixes: list): """ Split the dataset into training set and validation set. :return: train_item_names, valid_item_names """ prefixes = {str(pr): 1 for pr in prefixes} valid_item_names = {} # Add prefixes that specified speaker index and matches exactly item name to test set for prefix in deepcopy(prefixes): if prefix in self.item_names: valid_item_names[prefix] = 1 prefixes.pop(prefix) # Add prefixes that exactly matches item name without speaker id to test set for prefix in deepcopy(prefixes): matched = False for name in self.item_names: if name.split(':')[-1] == prefix: valid_item_names[name] = 1 matched = True if matched: prefixes.pop(prefix) # Add names with one of the remaining prefixes to test set for prefix in deepcopy(prefixes): matched = False for name in self.item_names: if name.startswith(prefix): valid_item_names[name] = 1 matched = True if matched: prefixes.pop(prefix) for prefix in deepcopy(prefixes): matched = False for name in self.item_names: if name.split(':')[-1].startswith(prefix): valid_item_names[name] = 1 matched = True if matched: prefixes.pop(prefix) if len(prefixes) != 0: warnings.warn( f'The following rules in test_prefixes have no matching names in the dataset: {", ".join(prefixes.keys())}', category=UserWarning ) warnings.filterwarnings('default') valid_item_names = list(valid_item_names.keys()) assert len(valid_item_names) > 0, 'Validation set is empty!' train_item_names = [x for x in self.item_names if x not in set(valid_item_names)] assert len(train_item_names) > 0, 'Training set is empty!' return train_item_names, valid_item_names @property def train_item_names(self): return self._train_item_names @property def valid_item_names(self): return self._valid_item_names def meta_data_iterator(self, prefix): if prefix == 'train': item_names = self.train_item_names else: item_names = self.valid_item_names for item_name in item_names: meta_data = self.items[item_name] yield item_name, meta_data def process(self): # load each dataset test_prefixes = [] for ds_id, dataset in enumerate(self.datasets): items = self.load_meta_data( pathlib.Path(dataset['raw_data_dir']), ds_id=ds_id, spk=dataset['speaker'], lang=dataset['language'] ) self.items.update(items) test_prefixes.extend( f'{ds_id}:{prefix}' for prefix in dataset.get('test_prefixes', []) ) self.item_names = sorted(list(self.items.keys())) self._train_item_names, self._valid_item_names = self.split_train_valid_set(test_prefixes) if self.binarization_args['shuffle']: random.shuffle(self.item_names) self.binary_data_dir.mkdir(parents=True, exist_ok=True) # Copy spk_map, lang_map and dictionary to binary data dir spk_map_fn = self.binary_data_dir / 'spk_map.json' with open(spk_map_fn, 'w', encoding='utf-8') as f: json.dump(self.spk_map, f, ensure_ascii=False) lang_map_fn = self.binary_data_dir / 'lang_map.json' with open(lang_map_fn, 'w', encoding='utf-8') as f: json.dump(self.lang_map, f, ensure_ascii=False) for lang, dict_path in hparams['dictionaries'].items(): shutil.copy(dict_path, self.binary_data_dir / f'dictionary-{lang}.txt') self.check_coverage() # Process valid set and train set try: self.process_dataset('valid') self.process_dataset( 'train', num_workers=int(self.binarization_args['num_workers']), apply_augmentation=any(args['enabled'] for args in self.augmentation_args.values()) ) except KeyboardInterrupt: exit(-1) def check_coverage(self): # Group by phonemes in the dictionary. ph_idx_required = set(range(1, len(self.phoneme_dictionary))) ph_idx_occurred = set() ph_idx_count_map = { idx: 0 for idx in ph_idx_required } # Load and count those phones that appear in the actual data for item_name in self.items: ph_idx_occurred.update(self.items[item_name]['ph_seq']) for idx in self.items[item_name]['ph_seq']: ph_idx_count_map[idx] += 1 ph_count_map = { self.phoneme_dictionary.decode_one(idx, scalar=False): count for idx, count in ph_idx_count_map.items() } def display_phoneme(phoneme): if isinstance(phoneme, tuple): return f'({", ".join(phoneme)})' return phoneme print('===== Phoneme Distribution Summary =====') keys = sorted(ph_count_map.keys(), key=lambda v: v[0] if isinstance(v, tuple) else v) for i, key in enumerate(keys): if i == len(ph_count_map) - 1: end = '\n' elif i % 10 == 9: end = ',\n' else: end = ', ' key_disp = display_phoneme(key) print(f'{key_disp}: {ph_count_map[key]}', end=end) # Draw graph. xs = [display_phoneme(k) for k in keys] ys = [ph_count_map[k] for k in keys] plt = distribution_to_figure( title='Phoneme Distribution Summary', x_label='Phoneme', y_label='Number of occurrences', items=xs, values=ys, rotate=len(self.dictionaries) > 1 ) filename = self.binary_data_dir / 'phoneme_distribution.jpg' plt.savefig(fname=filename, bbox_inches='tight', pad_inches=0.25) print(f'| save summary to \'{filename}\'') # Check unrecognizable or missing phonemes if ph_idx_occurred != ph_idx_required: missing_phones = sorted({ self.phoneme_dictionary.decode_one(idx, scalar=False) for idx in ph_idx_required.difference(ph_idx_occurred) }, key=lambda v: v[0] if isinstance(v, tuple) else v) raise BinarizationError( f'The following phonemes are not covered in transcriptions: {missing_phones}' ) def process_dataset(self, prefix, num_workers=0, apply_augmentation=False): args = [] builder = IndexedDatasetBuilder(self.binary_data_dir, prefix=prefix, allowed_attr=self.data_attrs) total_sec = {k: 0.0 for k in self.spk_map} total_raw_sec = {k: 0.0 for k in self.spk_map} extra_info = {'names': {}, 'ph_texts': {}, 'spk_ids': {}, 'spk_names': {}, 'lengths': {}} max_no = -1 for item_name, meta_data in self.meta_data_iterator(prefix): args.append([item_name, meta_data, self.binarization_args]) aug_map = self.arrange_data_augmentation(self.meta_data_iterator(prefix)) if apply_augmentation else {} def postprocess(_item): nonlocal total_sec, total_raw_sec, extra_info, max_no if _item is None: return item_no = builder.add_item(_item) max_no = max(max_no, item_no) for k, v in _item.items(): if isinstance(v, np.ndarray): if k not in extra_info: extra_info[k] = {} extra_info[k][item_no] = v.shape[0] extra_info['names'][item_no] = _item['name'].split(':', 1)[-1] extra_info['ph_texts'][item_no] = _item['ph_text'] extra_info['spk_ids'][item_no] = _item['spk_id'] extra_info['spk_names'][item_no] = _item['spk_name'] extra_info['lengths'][item_no] = _item['length'] total_raw_sec[_item['spk_name']] += _item['seconds'] total_sec[_item['spk_name']] += _item['seconds'] for task in aug_map.get(_item['name'], []): aug_item = task['func'](_item, **task['kwargs']) aug_item_no = builder.add_item(aug_item) max_no = max(max_no, aug_item_no) for k, v in aug_item.items(): if isinstance(v, np.ndarray): if k not in extra_info: extra_info[k] = {} extra_info[k][aug_item_no] = v.shape[0] extra_info['names'][aug_item_no] = aug_item['name'].split(':', 1)[-1] extra_info['ph_texts'][aug_item_no] = aug_item['ph_text'] extra_info['spk_ids'][aug_item_no] = aug_item['spk_id'] extra_info['spk_names'][aug_item_no] = aug_item['spk_name'] extra_info['lengths'][aug_item_no] = aug_item['length'] total_sec[aug_item['spk_name']] += aug_item['seconds'] try: if num_workers > 0: # code for parallel processing for item in tqdm( chunked_multiprocess_run(self.process_item, args, num_workers=num_workers), total=len(list(self.meta_data_iterator(prefix))) ): postprocess(item) else: # code for single cpu processing for a in tqdm(args): item = self.process_item(*a) postprocess(item) for k in extra_info: assert set(extra_info[k]) == set(range(max_no + 1)), f'Item numbering is not consecutive.' extra_info[k] = list(map(lambda x: x[1], sorted(extra_info[k].items(), key=lambda x: x[0]))) except KeyboardInterrupt: builder.finalize() raise builder.finalize() if prefix == "train": extra_info.pop("names") extra_info.pop('ph_texts') extra_info.pop("spk_names") with open(self.binary_data_dir / f"{prefix}.meta", "wb") as f: # noinspection PyTypeChecker pickle.dump(extra_info, f) if apply_augmentation: print(f"| {prefix} total duration (before augmentation): {sum(total_raw_sec.values()):.2f}s") print( f"| {prefix} respective duration (before augmentation): " + ', '.join(f'{k}={v:.2f}s' for k, v in total_raw_sec.items()) ) print( f"| {prefix} total duration (after augmentation): " f"{sum(total_sec.values()):.2f}s ({sum(total_sec.values()) / sum(total_raw_sec.values()):.2f}x)" ) print( f"| {prefix} respective duration (after augmentation): " + ', '.join(f'{k}={v:.2f}s' for k, v in total_sec.items()) ) else: print(f"| {prefix} total duration: {sum(total_raw_sec.values()):.2f}s") print(f"| {prefix} respective duration: " + ', '.join(f'{k}={v:.2f}s' for k, v in total_raw_sec.items())) def arrange_data_augmentation(self, data_iterator): """ Code for all types of data augmentation should be added here. """ raise NotImplementedError() def process_item(self, item_name, meta_data, binarization_args): raise NotImplementedError()