File size: 4,252 Bytes
7934b29 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# 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.
"""
This script is to compute speaker-level statistics, such as pitch mean & standard deviation, for a given
TTS training manifest.
This script should be run after extract_sup_data.py as it uses the precomputed supplemental features.
$ python <nemo_root_path>/scripts/dataset_processing/tts/compute_speaker_stats.py \
--manifest_path=<data_root_path>/fastpitch_manifest.json \
--sup_data_path=<data_root_path>/sup_data \
--pitch_stats_path=<data_root_path>/pitch_stats.json
"""
import argparse
import json
import os
from collections import defaultdict
from pathlib import Path
from typing import List, Tuple
import torch
from tqdm import tqdm
from nemo.collections.asr.parts.utils.manifest_utils import read_manifest
from nemo.collections.tts.parts.utils.tts_dataset_utils import get_base_dir, get_sup_data_file_path
from nemo.collections.tts.torch.tts_data_types import Pitch
from nemo.utils import logging
def get_args():
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter, description="Compute speaker level pitch statistics.",
)
parser.add_argument(
"--manifest_path", required=True, type=Path, help="Path to training manifest.",
)
parser.add_argument(
"--sup_data_path", default=Path("sup_data"), type=Path, help="Path to base directory with supplementary data.",
)
parser.add_argument(
"--pitch_stats_path",
default=Path("pitch_stats.json"),
type=Path,
help="Path to output JSON file with speaker pitch statistics.",
)
args = parser.parse_args()
return args
def _compute_stats(values: List[torch.Tensor]) -> Tuple[float, float]:
values_tensor = torch.cat(values, dim=0)
mean = values_tensor.mean().item()
std = values_tensor.std(dim=0).item()
return mean, std
def main():
args = get_args()
manifest_path = args.manifest_path
sup_data_path = args.sup_data_path
pitch_stats_path = args.pitch_stats_path
pitch_data_path = Path(os.path.join(sup_data_path, Pitch.name))
if not os.path.exists(pitch_data_path):
raise ValueError(
f"Pitch directory {pitch_data_path} does not exist. Make sure 'sup_data_path' is correct "
f"and that you have computed the pitch using extract_sup_data.py"
)
entries = read_manifest(manifest_path)
audio_paths = [entry["audio_filepath"] for entry in entries]
base_dir = get_base_dir(audio_paths)
global_pitch_values = []
speaker_pitch_values = defaultdict(list)
for entry in tqdm(entries):
pitch_path = get_sup_data_file_path(entry, base_dir, pitch_data_path)
if not os.path.exists(pitch_path):
logging.warning(f"Unable to find pitch file for {entry}")
continue
pitch = torch.load(pitch_path)
# Filter out non-speech frames
pitch = pitch[pitch != 0]
global_pitch_values.append(pitch)
if "speaker" in entry:
speaker_id = entry["speaker"]
speaker_pitch_values[speaker_id].append(pitch)
global_pitch_mean, global_pitch_std = _compute_stats(global_pitch_values)
pitch_stats = {"default": {"pitch_mean": global_pitch_mean, "pitch_std": global_pitch_std}}
for speaker_id, pitch_values in speaker_pitch_values.items():
pitch_mean, pitch_std = _compute_stats(pitch_values)
pitch_stats[speaker_id] = {"pitch_mean": pitch_mean, "pitch_std": pitch_std}
with open(pitch_stats_path, 'w', encoding="utf-8") as stats_f:
json.dump(pitch_stats, stats_f, indent=4)
if __name__ == "__main__":
main()
|