process_pipeline.py

'''
    file -> temporary_dict -> processed_input -> batch
'''
from utils.hparams import hparams
from network.vocoders.base_vocoder import VOCODERS
import numpy as np
import traceback
from pathlib import Path
from .data_gen_utils import get_pitch_parselmouth,get_pitch_crepe
from .base_binarizer import BinarizationError
import torch
import utils

class File2Batch:
    '''
        pipeline: file -> temporary_dict -> processed_input -> batch
    '''

    @staticmethod
    def file2temporary_dict():
        '''
            read from file, store data in temporary dicts
        '''
        raw_data_dir = Path(hparams['raw_data_dir'])
        # meta_midi = json.load(open(os.path.join(raw_data_dir, 'meta.json')))   # [list of dict]
        
        # if hparams['perform_enhance'] and not hparams['infer']:
        #     vocoder=get_vocoder_cls(hparams)()
        #     raw_files = list(raw_data_dir.rglob(f"*.wav"))
        #     dic=[]
        #     time_step = hparams['hop_size'] / hparams['audio_sample_rate']
        #     f0_min = hparams['f0_min']
        #     f0_max = hparams['f0_max']
        #     for file in raw_files:
        #         y, sr = librosa.load(file, sr=hparams['audio_sample_rate'])
        #         f0 = parselmouth.Sound(y, hparams['audio_sample_rate']).to_pitch_ac(
        #         time_step=time_step , voicing_threshold=0.6,
        #         pitch_floor=f0_min, pitch_ceiling=f0_max).selected_array['frequency']
        #         f0_mean=np.mean(f0[f0>0])
        #         dic.append(f0_mean)
        #     for idx in np.where(dic>np.percentile(dic, 80))[0]:
        #         file=raw_files[idx]
        #         wav,mel=vocoder.wav2spec(str(file))
        #         f0,_=get_pitch_parselmouth(wav,mel,hparams)
        #         f0[f0>0]=f0[f0>0]*(2**(2/12))
        #         wav_pred=vocoder.spec2wav(torch.FloatTensor(mel),f0=torch.FloatTensor(f0))
        #         sf.write(file.with_name(file.name[:-4]+'_high.wav'), wav_pred, 24000, 'PCM_16')
        utterance_labels =[]
        utterance_labels.extend(list(raw_data_dir.rglob(f"*.wav")))
        utterance_labels.extend(list(raw_data_dir.rglob(f"*.ogg")))
        #open(os.path.join(raw_data_dir, 'transcriptions.txt'), encoding='utf-8').readlines()

        all_temp_dict = {}
        for utterance_label in utterance_labels:
            #song_info = utterance_label.split('|')
            item_name =str(utterance_label)#raw_item_name = song_info[0]
            # print(item_name)
            temp_dict = {}
            temp_dict['wav_fn'] =str(utterance_label)#f'{raw_data_dir}/wavs/{item_name}.wav'
            # temp_dict['txt'] = song_info[1]

            # temp_dict['ph'] = song_info[2]
            # # self.item2wdb[item_name] = list(np.nonzero([1 if x in ALL_YUNMU + ['AP', 'SP'] else 0 for x in song_info[2].split()])[0])
            # temp_dict['word_boundary'] = np.array([1 if x in ALL_YUNMU + ['AP', 'SP'] else 0 for x in song_info[2].split()])
            # temp_dict['ph_durs'] = [float(x) for x in song_info[5].split(" ")]

            # temp_dict['pitch_midi'] = np.array([note_to_midi(x.split("/")[0]) if x != 'rest' else 0
            #                        for x in song_info[3].split(" ")])
            # temp_dict['midi_dur'] = np.array([float(x) for x in song_info[4].split(" ")])
            # temp_dict['is_slur'] = np.array([int(x) for x in song_info[6].split(" ")])
            temp_dict['spk_id'] = hparams['speaker_id']
            # assert temp_dict['pitch_midi'].shape == temp_dict['midi_dur'].shape == temp_dict['is_slur'].shape, \
                # (temp_dict['pitch_midi'].shape, temp_dict['midi_dur'].shape, temp_dict['is_slur'].shape)

            all_temp_dict[item_name] = temp_dict

        return all_temp_dict
    
    @staticmethod
    def temporary_dict2processed_input(item_name, temp_dict, encoder, binarization_args):
        '''
            process data in temporary_dicts
        '''
        def get_pitch(wav, mel):
            # get ground truth f0 by self.get_pitch_algorithm
            if hparams['use_crepe']:
                gt_f0, gt_pitch_coarse = get_pitch_crepe(wav, mel, hparams)
            else:
                gt_f0, gt_pitch_coarse = get_pitch_parselmouth(wav, mel, hparams)
            if sum(gt_f0) == 0:
                raise BinarizationError("Empty **gt** f0")
            processed_input['f0'] = gt_f0
            processed_input['pitch'] = gt_pitch_coarse
    
        def get_align(meta_data, mel, phone_encoded, hop_size=hparams['hop_size'], audio_sample_rate=hparams['audio_sample_rate']):
            mel2ph = np.zeros([mel.shape[0]], int)
            start_frame=0
            ph_durs = mel.shape[0]/phone_encoded.shape[0]
            if hparams['debug']:
                print(mel.shape,phone_encoded.shape,mel.shape[0]/phone_encoded.shape[0])
            for i_ph in range(phone_encoded.shape[0]):
                
                end_frame = int(i_ph*ph_durs +ph_durs+ 0.5)
                mel2ph[start_frame:end_frame+1] = i_ph + 1
                start_frame = end_frame+1

            processed_input['mel2ph'] = mel2ph

        if hparams['vocoder'] in VOCODERS:
            wav, mel = VOCODERS[hparams['vocoder']].wav2spec(temp_dict['wav_fn'])
        else:
            wav, mel = VOCODERS[hparams['vocoder'].split('.')[-1]].wav2spec(temp_dict['wav_fn'])
        processed_input = {
            'item_name': item_name, 'mel': mel, 'wav': wav,
            'sec': len(wav) / hparams['audio_sample_rate'], 'len': mel.shape[0]
        }
        processed_input = {**temp_dict, **processed_input} # merge two dicts
        processed_input['spec_min']=np.min(mel,axis=0)
        processed_input['spec_max']=np.max(mel,axis=0)
        #(processed_input['spec_min'].shape)
        try:
            if binarization_args['with_f0']:
                get_pitch(wav, mel)
            if binarization_args['with_hubert']:
                try:
                    hubert_encoded = processed_input['hubert'] = encoder.encode(temp_dict['wav_fn'])
                except:
                    traceback.print_exc()
                    raise Exception(f"hubert encode error")
                if binarization_args['with_align']:
                    get_align(temp_dict, mel, hubert_encoded)
        except Exception as e:
            print(f"| Skip item ({e}). item_name: {item_name}, wav_fn: {temp_dict['wav_fn']}")
            return None
        return processed_input

    @staticmethod
    def processed_input2batch(samples):
        '''
            Args:
                samples: one batch of processed_input
            NOTE:
                the batch size is controlled by hparams['max_sentences']
        '''
        if len(samples) == 0:
            return {}
        id = torch.LongTensor([s['id'] for s in samples])
        item_names = [s['item_name'] for s in samples]
        #text = [s['text'] for s in samples]
        #txt_tokens = utils.collate_1d([s['txt_token'] for s in samples], 0)
        hubert = utils.collate_2d([s['hubert'] for s in samples], 0.0)
        f0 = utils.collate_1d([s['f0'] for s in samples], 0.0)
        pitch = utils.collate_1d([s['pitch'] for s in samples])
        uv = utils.collate_1d([s['uv'] for s in samples])
        energy = utils.collate_1d([s['energy'] for s in samples], 0.0)
        mel2ph = utils.collate_1d([s['mel2ph'] for s in samples], 0.0) \
            if samples[0]['mel2ph'] is not None else None
        mels = utils.collate_2d([s['mel'] for s in samples], 0.0)
        #txt_lengths = torch.LongTensor([s['txt_token'].numel() for s in samples])
        hubert_lengths = torch.LongTensor([s['hubert'].shape[0] for s in samples])
        mel_lengths = torch.LongTensor([s['mel'].shape[0] for s in samples])

        batch = {
            'id': id,
            'item_name': item_names,
            'nsamples': len(samples),
            # 'text': text,
            # 'txt_tokens': txt_tokens,
            # 'txt_lengths': txt_lengths,
            'hubert':hubert,
            'mels': mels,
            'mel_lengths': mel_lengths,
            'mel2ph': mel2ph,
            'energy': energy,
            'pitch': pitch,
            'f0': f0,
            'uv': uv,
        }
        #========not used=================
        # if hparams['use_spk_embed']:
        #     spk_embed = torch.stack([s['spk_embed'] for s in samples])
        #     batch['spk_embed'] = spk_embed
        # if hparams['use_spk_id']:
        #     spk_ids = torch.LongTensor([s['spk_id'] for s in samples])
        #     batch['spk_ids'] = spk_ids
        # if hparams['pitch_type'] == 'cwt':
        #     cwt_spec = utils.collate_2d([s['cwt_spec'] for s in samples])
        #     f0_mean = torch.Tensor([s['f0_mean'] for s in samples])
        #     f0_std = torch.Tensor([s['f0_std'] for s in samples])
        #     batch.update({'cwt_spec': cwt_spec, 'f0_mean': f0_mean, 'f0_std': f0_std})
        # elif hparams['pitch_type'] == 'ph':
        #     batch['f0'] = utils.collate_1d([s['f0_ph'] for s in samples])

        # batch['pitch_midi'] = utils.collate_1d([s['pitch_midi'] for s in samples], 0)
        # batch['midi_dur'] = utils.collate_1d([s['midi_dur'] for s in samples], 0)
        # batch['is_slur'] = utils.collate_1d([s['is_slur'] for s in samples], 0)
        # batch['word_boundary'] = utils.collate_1d([s['word_boundary'] for s in samples], 0)

        return batch