|
- # Copyright 2021 Huawei Technologies Co., Ltd
- #
- # Licensed under the Apache License, Version 2.0 (the "License");
- # you may not use this file except in compliance with the License.
- # You may obtain a copy of the License at
- #
- # http://www.apache.org/licenses/LICENSE-2.0
- #
- # Unless required by applicable law or agreed to in writing, software
- # distributed under the License is distributed on an "AS IS" BASIS,
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- # See the License for the specific language governing permissions and
- # limitations under the License.
- # ============================================================================
- """
- Evaluation.
- """
- import os
- from os.path import join
- import argparse
- import glob
- import math
- import audio
- import numpy as np
- from scipy.io import wavfile
- from hparams import hparams, hparams_debug_string
- from tqdm import tqdm
- from mindspore import context, Tensor
- from mindspore.train.serialization import load_checkpoint, load_param_into_net
- import mindspore.dataset.engine as de
- from nnmnkwii import preprocessing as P
- from nnmnkwii.datasets import FileSourceDataset
- from wavenet_vocoder import WaveNet
- from wavenet_vocoder.util import is_mulaw_quantize, is_mulaw, is_scalar_input
- from src.dataset import RawAudioDataSource, MelSpecDataSource, DualDataset
-
- parser = argparse.ArgumentParser(description='TTS training')
- parser.add_argument('--data_path', type=str, required=True, default='',
- help='Directory contains preprocessed features.')
- parser.add_argument('--preset', type=str, required=True, default='', help='Path of preset parameters (json).')
- parser.add_argument('--pretrain_ckpt', type=str, default='', help='Pretrained checkpoint path')
- parser.add_argument('--is_numpy', action="store_true", default=False, help='Using numpy for inference or not')
- parser.add_argument('--output_path', type=str, default='./out_wave/', help='Path to save generated audios')
- parser.add_argument('--platform', type=str, default='GPU', choices=('Ascend', 'GPU', 'CPU'),
- help='run platform, support Ascend, GPU and CPU. Default: GPU')
- parser.add_argument('--speaker_id', type=str, default='',
- help=' Use specific speaker of data in case for multi-speaker datasets.')
- args = parser.parse_args()
-
-
- def get_data_loader(hparam, data_dir):
- """
- test data loader
- """
- wav_paths = glob.glob(os.path.join(data_dir, "*-wave.npy"))
- if wav_paths:
- X = FileSourceDataset(RawAudioDataSource(data_dir,
- hop_size=audio.get_hop_size(),
- max_steps=None, cin_pad=hparam.cin_pad))
- else:
- X = None
- C = FileSourceDataset(MelSpecDataSource(data_dir,
- hop_size=audio.get_hop_size(),
- max_steps=None, cin_pad=hparam.cin_pad))
-
- length_x = np.array(C.file_data_source.lengths)
- if C[0].shape[-1] != hparam.cin_channels:
- raise RuntimeError("Invalid cin_channnels {}. Expected to be {}.".format(hparam.cin_channels, C[0].shape[-1]))
-
- dataset = DualDataset(X, C, length_x, batch_size=hparam.batch_size, hparams=hparam)
-
- data_loader = de.GeneratorDataset(dataset, ["x_batch", "y_batch", "c_batch", "g_batch", "input_lengths", "mask"])
-
- return data_loader, dataset
-
-
- def batch_wavegen(hparam, net, c_input=None, g_input=None, tqdm_=None, is_numpy=True):
- """
- generate audio
- """
- assert c_input is not None
- B = c_input.shape[0]
- net.set_train(False)
-
- n_frames = c_input.shape[-1]
- y_hat_list = []
- chunk_wise = 16 + 2 * hparam.cin_pad
- for k in range(math.ceil(n_frames / chunk_wise)):
- start = k * chunk_wise
- end = min(n_frames, (k + 1) * chunk_wise)
- lens = (end - start - hparam.cin_pad * 2) * audio.get_hop_size()
-
- y_hat = net.incremental_forward(c_=c_input[:, :, start: end], g=g_input, T=lens, tqdm=tqdm_, softmax=True,
- quantize=True,
- log_scale_min=hparam.log_scale_min, is_numpy=is_numpy)
- y_hat_list.append(y_hat)
-
- y_hat = np.concatenate(y_hat_list, axis=2)
-
- if is_mulaw_quantize(hparam.input_type):
- # needs to be float since mulaw_inv returns in range of [-1, 1]
- y_hat = np.reshape(np.argmax(y_hat, 1), (B, -1))
- y_hat = y_hat.astype(np.float32)
- for k in range(B):
- y_hat[k] = P.inv_mulaw_quantize(y_hat[k], hparam.quantize_channels - 1)
- elif is_mulaw(hparam.input_type):
- y_hat = np.reshape(y_hat, (B, -1))
- for k in range(B):
- y_hat[k] = P.inv_mulaw(y_hat[k], hparam.quantize_channels - 1)
- else:
- y_hat = np.reshape(y_hat, (B, -1))
-
- if hparam.postprocess is not None and hparam.postprocess not in ["", "none"]:
- for k in range(B):
- y_hat[k] = getattr(audio, hparam.postprocess)(y_hat[k])
-
- if hparam.global_gain_scale > 0:
- for k in range(B):
- y_hat[k] /= hparam.global_gain_scale
-
- return y_hat
-
-
- def to_int16(x_):
- """
- convert datatype to int16
- """
- if x_.dtype == np.int16:
- return x_
- assert x_.dtype == np.float32
- assert x_.min() >= -1 and x_.max() <= 1.0
- return (x_ * 32767).astype(np.int16)
-
-
- def get_reference_file(hparam, dataset_source, idx):
- """
- get reference files
- """
- reference_files = []
- reference_feats = []
- for _ in range(hparam.batch_size):
- if hasattr(dataset_source, "X"):
- reference_files.append(dataset_source.X.collected_files[idx][0])
- else:
- pass
- if hasattr(dataset_source, "Mel"):
- reference_feats.append(dataset_source.Mel.collected_files[idx][0])
- else:
- reference_feats.append(dataset_source.collected_files[idx][0])
- idx += 1
- return reference_files, reference_feats, idx
-
-
- def get_saved_audio_name(has_ref_file_, ref_file, ref_feat, g_fp):
- """get path to save reference audio"""
- if has_ref_file_:
- target_audio_path = ref_file
- name = os.path.splitext(os.path.basename(target_audio_path))[0].replace("-wave", "")
- else:
- target_feat_path = ref_feat
- name = os.path.splitext(os.path.basename(target_feat_path))[0].replace("-feats", "")
- # Paths
- if g_fp is None:
- dst_wav_path_ = join(args.output_path, "{}_gen.wav".format(name))
- target_wav_path_ = join(args.output_path, "{}_ref.wav".format(name))
- else:
- dst_wav_path_ = join(args.output_path, "speaker{}_{}_gen.wav".format(g, name))
- target_wav_path_ = join(args.output_path, "speaker{}_{}_ref.wav".format(g, name))
- return dst_wav_path_, target_wav_path_
-
-
- def save_ref_audio(hparam, ref, length, target_wav_path_):
- """
- save reference audio
- """
- if is_mulaw_quantize(hparam.input_type):
- ref = np.reshape(np.argmax(ref, 0), (-1))[:length]
- ref = ref.astype(np.float32)
- else:
- ref = np.reshape(ref, (-1))[:length]
-
- if is_mulaw_quantize(hparam.input_type):
- ref = P.inv_mulaw_quantize(ref, hparam.quantize_channels - 1)
- elif is_mulaw(hparam.input_type):
- ref = P.inv_mulaw(ref, hparam.quantize_channels - 1)
- if hparam.postprocess is not None and hparam.postprocess not in ["", "none"]:
- ref = getattr(audio, hparam.postprocess)(ref)
- if hparam.global_gain_scale > 0:
- ref /= hparam.global_gain_scale
-
- ref = np.clip(ref, -1.0, 1.0)
-
- wavfile.write(target_wav_path_, hparam.sample_rate, to_int16(ref))
-
-
- if __name__ == '__main__':
-
- if args.platform != 'Ascend':
- context.set_context(mode=0, device_target=args.platform, save_graphs=False)
- else:
- device_id = int(os.getenv("DEVICE_ID"))
- context.set_context(mode=1, device_target=args.platform, device_id=device_id)
-
- speaker_id = int(args.speaker_id) if args.speaker_id != '' else None
- if args.preset is not None:
- with open(args.preset) as f:
- hparams.parse_json(f.read())
-
- assert hparams.name == "wavenet_vocoder"
- print(hparams_debug_string())
-
- fs = hparams.sample_rate
- hparams.batch_size = 10
- hparams.max_time_sec = None
- hparams.max_time_steps = None
- data_loaders, source_dataset = get_data_loader(hparam=hparams, data_dir=args.data_path)
-
- upsample_params = hparams.upsample_params
- upsample_params["cin_channels"] = hparams.cin_channels
- upsample_params["cin_pad"] = hparams.cin_pad
- model = WaveNet(
- out_channels=hparams.out_channels,
- layers=hparams.layers,
- stacks=hparams.stacks,
- residual_channels=hparams.residual_channels,
- gate_channels=hparams.gate_channels,
- skip_out_channels=hparams.skip_out_channels,
- cin_channels=hparams.cin_channels,
- gin_channels=hparams.gin_channels,
- n_speakers=hparams.n_speakers,
- dropout=hparams.dropout,
- kernel_size=hparams.kernel_size,
- cin_pad=hparams.cin_pad,
- upsample_conditional_features=hparams.upsample_conditional_features,
- upsample_params=upsample_params,
- scalar_input=is_scalar_input(hparams.input_type),
- output_distribution=hparams.output_distribution,
- )
-
- param_dict = load_checkpoint(args.pretrain_ckpt)
- load_param_into_net(model, param_dict)
- print('Successfully loading the pre-trained model')
-
- os.makedirs(args.output_path, exist_ok=True)
- cin_pad = hparams.cin_pad
-
- file_idx = 0
- for data in data_loaders.create_dict_iterator():
- x, y, c, g, input_lengths = data['x_batch'], data['y_batch'], data['c_batch'], data['g_batch'], data[
- 'input_lengths']
- if cin_pad > 0:
- c = c.asnumpy()
- c = np.pad(c, pad_width=(cin_pad, cin_pad), mode="edge")
- c = Tensor(c)
-
- ref_files, ref_feats, file_idx = get_reference_file(hparams, source_dataset, file_idx)
- # Generate
- y_hats = batch_wavegen(hparams, model, data['c_batch'], tqdm_=tqdm, is_numpy=args.is_numpy)
- x = x.asnumpy()
- input_lengths = input_lengths.asnumpy()
- # Save each utt.
- has_ref_file = bool(ref_files)
- for i, (ref_, gen_, length_) in enumerate(zip(x, y_hats, input_lengths)):
- dst_wav_path, target_wav_path = get_saved_audio_name(has_ref_file_=has_ref_file, ref_file=ref_files[i],
- ref_feat=ref_feats[i], g_fp=g)
- save_ref_audio(hparams, ref_, length_, target_wav_path)
-
- gen = gen_[:length_]
- gen = np.clip(gen, -1.0, 1.0)
- wavfile.write(dst_wav_path, hparams.sample_rate, to_int16(gen))
|