wehub-resource-sync
387 lines
16 KiB
Python
387 lines
16 KiB
Python
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()
|