Comparative-Analysis-of-Speech-Synthesis-Models
/
TensorFlowTTS
/examples
/tacotron2
/extract_postnets.py
# -*- coding: utf-8 -*- | |
# Copyright 2020 Minh Nguyen (@dathudeptrai) | |
# | |
# 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. | |
"""Extract durations based-on tacotron-2 alignments for FastSpeech.""" | |
import argparse | |
import logging | |
import os | |
from numba import jit | |
import sys | |
sys.path.append(".") | |
import matplotlib.pyplot as plt | |
import numpy as np | |
import tensorflow as tf | |
import yaml | |
from tqdm import tqdm | |
from examples.tacotron2.tacotron_dataset import CharactorMelDataset | |
from tensorflow_tts.configs import Tacotron2Config | |
from tensorflow_tts.models import TFTacotron2 | |
def get_duration_from_alignment(alignment): | |
D = np.array([0 for _ in range(np.shape(alignment)[0])]) | |
for i in range(np.shape(alignment)[1]): | |
max_index = list(alignment[:, i]).index(alignment[:, i].max()) | |
D[max_index] = D[max_index] + 1 | |
return D | |
def main(): | |
"""Running extract tacotron-2 durations.""" | |
parser = argparse.ArgumentParser( | |
description="Extract durations from charactor with trained Tacotron-2 " | |
"(See detail in tensorflow_tts/example/tacotron-2/extract_duration.py)." | |
) | |
parser.add_argument( | |
"--rootdir", | |
default=None, | |
type=str, | |
required=True, | |
help="directory including ids/durations files.", | |
) | |
parser.add_argument( | |
"--outdir", type=str, required=True, help="directory to save generated mels." | |
) | |
parser.add_argument( | |
"--checkpoint", type=str, required=True, help="checkpoint file to be loaded." | |
) | |
parser.add_argument( | |
"--use-norm", default=1, type=int, help="usr norm-mels for train or raw." | |
) | |
parser.add_argument("--batch-size", default=32, type=int, help="batch size.") | |
parser.add_argument("--win-front", default=3, type=int, help="win-front.") | |
parser.add_argument("--win-back", default=3, type=int, help="win-front.") | |
parser.add_argument( | |
"--use-window-mask", default=1, type=int, help="toggle window masking." | |
) | |
parser.add_argument("--save-alignment", default=0, type=int, help="save-alignment.") | |
parser.add_argument( | |
"--config", | |
default=None, | |
type=str, | |
required=True, | |
help="yaml format configuration file. if not explicitly provided, " | |
"it will be searched in the checkpoint directory. (default=None)", | |
) | |
parser.add_argument( | |
"--verbose", | |
type=int, | |
default=1, | |
help="logging level. higher is more logging. (default=1)", | |
) | |
args = parser.parse_args() | |
# set logger | |
if args.verbose > 1: | |
logging.basicConfig( | |
level=logging.DEBUG, | |
format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s", | |
) | |
elif args.verbose > 0: | |
logging.basicConfig( | |
level=logging.INFO, | |
format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s", | |
) | |
else: | |
logging.basicConfig( | |
level=logging.WARN, | |
format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s", | |
) | |
logging.warning("Skip DEBUG/INFO messages") | |
# check directory existence | |
if not os.path.exists(args.outdir): | |
os.makedirs(args.outdir) | |
# load config | |
with open(args.config) as f: | |
config = yaml.load(f, Loader=yaml.Loader) | |
config.update(vars(args)) | |
if config["format"] == "npy": | |
char_query = "*-ids.npy" | |
mel_query = "*-raw-feats.npy" if args.use_norm is False else "*-norm-feats.npy" | |
char_load_fn = np.load | |
mel_load_fn = np.load | |
else: | |
raise ValueError("Only npy is supported.") | |
# define data-loader | |
dataset = CharactorMelDataset( | |
dataset=config["tacotron2_params"]["dataset"], | |
root_dir=args.rootdir, | |
charactor_query=char_query, | |
mel_query=mel_query, | |
charactor_load_fn=char_load_fn, | |
mel_load_fn=mel_load_fn, | |
reduction_factor=config["tacotron2_params"]["reduction_factor"], | |
use_fixed_shapes=True, | |
) | |
dataset = dataset.create( | |
allow_cache=True, batch_size=args.batch_size, drop_remainder=False | |
) | |
# define model and load checkpoint | |
tacotron2 = TFTacotron2( | |
config=Tacotron2Config(**config["tacotron2_params"]), | |
name="tacotron2", | |
) | |
tacotron2._build() # build model to be able load_weights. | |
tacotron2.load_weights(args.checkpoint) | |
# apply tf.function for tacotron2. | |
tacotron2 = tf.function(tacotron2, experimental_relax_shapes=True) | |
for data in tqdm(dataset, desc="[Extract Postnets]"): | |
utt_ids = data["utt_ids"] | |
input_lengths = data["input_lengths"] | |
mel_lengths = data["mel_lengths"] | |
utt_ids = utt_ids.numpy() | |
real_mel_lengths = data["real_mel_lengths"] | |
mel_gt = data["mel_gts"] | |
del data["real_mel_lengths"] | |
# tacotron2 inference. | |
mel_outputs, post_mel_outputs, stop_outputs, alignment_historys = tacotron2( | |
**data, | |
use_window_mask=args.use_window_mask, | |
win_front=args.win_front, | |
win_back=args.win_back, | |
training=True, | |
) | |
# convert to numpy | |
alignment_historys = alignment_historys.numpy() | |
post_mel_outputs = post_mel_outputs.numpy() | |
mel_gt = mel_gt.numpy() | |
outdpost = os.path.join(args.outdir, "postnets") | |
if not os.path.exists(outdpost): | |
os.makedirs(outdpost) | |
for i, alignment in enumerate(alignment_historys): | |
real_char_length = input_lengths[i].numpy() | |
real_mel_length = real_mel_lengths[i].numpy() | |
alignment_mel_length = int( | |
np.ceil( | |
real_mel_length / config["tacotron2_params"]["reduction_factor"] | |
) | |
) | |
alignment = alignment[:real_char_length, :alignment_mel_length] | |
d = get_duration_from_alignment(alignment) # [max_char_len] | |
d = d * config["tacotron2_params"]["reduction_factor"] | |
assert ( | |
np.sum(d) >= real_mel_length | |
), f"{d}, {np.sum(d)}, {alignment_mel_length}, {real_mel_length}" | |
if np.sum(d) > real_mel_length: | |
rest = np.sum(d) - real_mel_length | |
# print(d, np.sum(d), real_mel_length) | |
if d[-1] > rest: | |
d[-1] -= rest | |
elif d[0] > rest: | |
d[0] -= rest | |
else: | |
d[-1] -= rest // 2 | |
d[0] -= rest - rest // 2 | |
assert d[-1] >= 0 and d[0] >= 0, f"{d}, {np.sum(d)}, {real_mel_length}" | |
saved_name = utt_ids[i].decode("utf-8") | |
# check a length compatible | |
assert ( | |
len(d) == real_char_length | |
), f"different between len_char and len_durations, {len(d)} and {real_char_length}" | |
assert ( | |
np.sum(d) == real_mel_length | |
), f"different between sum_durations and len_mel, {np.sum(d)} and {real_mel_length}" | |
# save D to folder. | |
np.save( | |
os.path.join(outdpost, f"{saved_name}-postnet.npy"), | |
post_mel_outputs[i][:][:real_mel_length].astype(np.float32), | |
allow_pickle=False, | |
) | |
# save alignment to debug. | |
if args.save_alignment == 1: | |
figname = os.path.join(args.outdir, f"{saved_name}_alignment.png") | |
fig = plt.figure(figsize=(8, 6)) | |
ax = fig.add_subplot(111) | |
ax.set_title(f"Alignment of {saved_name}") | |
im = ax.imshow( | |
alignment, aspect="auto", origin="lower", interpolation="none" | |
) | |
fig.colorbar(im, ax=ax) | |
xlabel = "Decoder timestep" | |
plt.xlabel(xlabel) | |
plt.ylabel("Encoder timestep") | |
plt.tight_layout() | |
plt.savefig(figname) | |
plt.close() | |
if __name__ == "__main__": | |
main() | |