mirror of https://github.com/coqui-ai/TTS.git
70 lines
2.2 KiB
Python
70 lines
2.2 KiB
Python
|
|
import argparse
|
|
import os
|
|
|
|
from tqdm import tqdm
|
|
|
|
from TTS.config import load_config
|
|
from TTS.tts.datasets import load_meta_data
|
|
from TTS.tts.utils.speakers import SpeakerManager
|
|
|
|
|
|
parser = argparse.ArgumentParser(
|
|
description='Compute embedding vectors for each wav file in a dataset.'
|
|
)
|
|
parser.add_argument("model_path", type=str, help="Path to model checkpoint file.")
|
|
parser.add_argument(
|
|
"config_path",
|
|
type=str,
|
|
help="Path to model config file.",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"config_dataset_path",
|
|
type=str,
|
|
help="Path to dataset config file.",
|
|
)
|
|
parser.add_argument("output_path", type=str, help="path for output speakers.json and/or speakers.npy.")
|
|
parser.add_argument("--use_cuda", type=bool, help="flag to set cuda.", default=True)
|
|
parser.add_argument("--eval", type=bool, help="compute eval.", default=True)
|
|
|
|
args = parser.parse_args()
|
|
|
|
c_dataset = load_config(args.config_dataset_path)
|
|
|
|
meta_data_train, meta_data_eval = load_meta_data(c_dataset.datasets, eval_split=args.eval)
|
|
wav_files = meta_data_train + meta_data_eval
|
|
|
|
speaker_manager = SpeakerManager(encoder_model_path=args.model_path, encoder_config_path=args.config_path, use_cuda=args.use_cuda)
|
|
|
|
# compute speaker embeddings
|
|
speaker_mapping = {}
|
|
for idx, wav_file in enumerate(tqdm(wav_files)):
|
|
if isinstance(wav_file, list):
|
|
speaker_name = wav_file[2]
|
|
wav_file = wav_file[1]
|
|
else:
|
|
speaker_name = None
|
|
|
|
# extract the embedding
|
|
embedd = speaker_manager.compute_d_vector_from_clip(wav_file)
|
|
|
|
# create speaker_mapping if target dataset is defined
|
|
wav_file_name = os.path.basename(wav_file)
|
|
speaker_mapping[wav_file_name] = {}
|
|
speaker_mapping[wav_file_name]["name"] = speaker_name
|
|
speaker_mapping[wav_file_name]["embedding"] = embedd
|
|
|
|
if speaker_mapping:
|
|
# save speaker_mapping if target dataset is defined
|
|
if '.json' not in args.output_path:
|
|
mapping_file_path = os.path.join(args.output_path, "speakers.json")
|
|
else:
|
|
mapping_file_path = args.output_path
|
|
|
|
os.makedirs(os.path.dirname(mapping_file_path), exist_ok=True)
|
|
|
|
# pylint: disable=W0212
|
|
speaker_manager._save_json(mapping_file_path, speaker_mapping)
|
|
print("Speaker embeddings saved at:", mapping_file_path)
|