basics/base_binarizer.py

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 build_phoneme_list, locate_dictionary
from utils.plot import distribution_to_figure
from utils.text_encoder import TokenTextEncoder


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, data_dir=None, data_attrs=None):
        if data_dir is None:
            data_dir = hparams['raw_data_dir']
        if not isinstance(data_dir, list):
            data_dir = [data_dir]

        self.raw_data_dirs = [pathlib.Path(d) for d in data_dir]
        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 = None
        self.spk_ids = hparams['spk_ids']
        self.speakers = hparams['speakers']
        self.build_spk_map()

        self.items = {}
        self.item_names: list = None
        self._train_item_names: list = None
        self._valid_item_names: list = None

        self.phone_encoder = TokenTextEncoder(vocab_list=build_phoneme_list())
        self.timestep = hparams['hop_size'] / hparams['audio_sample_rate']

    def build_spk_map(self):
        assert isinstance(self.speakers, list), 'Speakers must be a list'
        assert len(self.speakers) == len(self.raw_data_dirs), \
            'Number of raw data dirs must equal number of speaker names!'
        if len(self.spk_ids) == 0:
            self.spk_ids = list(range(len(self.raw_data_dirs)))
        else:
            assert len(self.spk_ids) == len(self.raw_data_dirs), \
                'Length of explicitly given spk_ids must equal the number of raw datasets.'
        assert max(self.spk_ids) < hparams['num_spk'], \
            f'Index in spk_id sequence {self.spk_ids} is out of range. All values should be smaller than num_spk.'

        self.spk_map = {}
        for spk_name, spk_id in zip(self.speakers, self.spk_ids):
            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

        print("| spk_map: ", self.spk_map)

    def load_meta_data(self, raw_data_dir: pathlib.Path, ds_id, spk_id):
        raise NotImplementedError()

    def split_train_valid_set(self, item_names):
        """
        Split the dataset into training set and validation set.
        :return: train_item_names, valid_item_names
        """
        prefixes = {str(pr): 1 for pr in hparams['test_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 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 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 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 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 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
        for ds_id, spk_id, data_dir in zip(range(len(self.raw_data_dirs)), self.spk_ids, self.raw_data_dirs):
            self.load_meta_data(pathlib.Path(data_dir), ds_id=ds_id, spk_id=spk_id)
        self.item_names = sorted(list(self.items.keys()))
        self._train_item_names, self._valid_item_names = self.split_train_valid_set(self.item_names)

        if self.binarization_args['shuffle']:
            random.shuffle(self.item_names)

        self.binary_data_dir.mkdir(parents=True, exist_ok=True)

        # Copy spk_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)
        shutil.copy(locate_dictionary(), self.binary_data_dir / 'dictionary.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_required = set(build_phoneme_list())
        phoneme_map = {}
        for ph in ph_required:
            phoneme_map[ph] = 0
        ph_occurred = []

        # Load and count those phones that appear in the actual data
        for item_name in self.items:
            ph_occurred += self.items[item_name]['ph_seq']
            if len(ph_occurred) == 0:
                raise BinarizationError(f'Empty tokens in {item_name}.')
        for ph in ph_occurred:
            if ph not in ph_required:
                continue
            phoneme_map[ph] += 1
        ph_occurred = set(ph_occurred)

        print('===== Phoneme Distribution Summary =====')
        for i, key in enumerate(sorted(phoneme_map.keys())):
            if i == len(ph_required) - 1:
                end = '\n'
            elif i % 10 == 9:
                end = ',\n'
            else:
                end = ', '
            print(f'\'{key}\': {phoneme_map[key]}', end=end)

        # Draw graph.
        x = sorted(phoneme_map.keys())
        values = [phoneme_map[k] for k in x]
        plt = distribution_to_figure(
            title='Phoneme Distribution Summary',
            x_label='Phoneme', y_label='Number of occurrences',
            items=x, values=values
        )
        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_occurred != ph_required:
            unrecognizable_phones = ph_occurred.difference(ph_required)
            missing_phones = ph_required.difference(ph_occurred)
            raise BinarizationError('transcriptions and dictionary mismatch.\n'
                                    f' (+) {sorted(unrecognizable_phones)}\n'
                                    f' (-) {sorted(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': {}, '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['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['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("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()