Spaces:
Paused
Paused
File size: 5,310 Bytes
2967cdb |
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 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 |
#!/usr/bin/env python3
# Copyright 2025 Xiaomi Corp. (authors: Han Zhu)
#
# See ../../../../LICENSE for clarification regarding multiple authors
#
# 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.
"""
Calculate UTMOS score with automatic Mean Opinion Score (MOS) prediction system
"""
import argparse
import logging
import os
from typing import List
import numpy as np
import torch
from tqdm import tqdm
from zipvoice.eval.models.utmos import UTMOS22Strong
from zipvoice.eval.utils import load_waveform
def get_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
description="Calculate UTMOS score using UTMOS22Strong model."
)
parser.add_argument(
"--wav-path",
type=str,
required=True,
help="Path to the directory containing evaluated speech files.",
)
parser.add_argument(
"--model-dir",
type=str,
required=True,
help="Local path of our evaluatioin model repository."
"Download from https://huggingface.co/k2-fsa/TTS_eval_models."
"Will use 'tts_eval_models/mos/utmos22_strong_step7459_v1.pt'"
" in this script",
)
parser.add_argument(
"--extension",
type=str,
default="wav",
help="Extension of the speech files. Default: wav",
)
return parser
class UTMOSScore:
"""Predicting UTMOS score for each audio clip."""
def __init__(self, model_path: str):
"""
Initializes the UTMOS score evaluator with the specified model.
Args:
model_path (str): Path of the UTMOS model checkpoint.
"""
self.sample_rate = 16000
self.device = (
torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
)
logging.info(f"Using device: {self.device}")
# Initialize and load the model
self.model = UTMOS22Strong()
try:
state_dict = torch.load(
model_path, map_location=lambda storage, loc: storage
)
self.model.load_state_dict(state_dict)
except Exception as e:
logging.error(f"Failed to load model from {model_path}: {e}")
raise
self.model.to(self.device)
self.model.eval()
@torch.no_grad()
def score_files(self, wav_paths: List[str]) -> List[float]:
"""
Computes UTMOS scores for a list of audio files.
Args:
wav_paths (List[str]): List of paths to audio files.
Returns:
List[float]: List of UTMOS scores.
"""
scores = []
for wav_path in tqdm(wav_paths, desc="Scoring audio files"):
# Load and preprocess waveform
speech = load_waveform(wav_path, self.sample_rate, device=self.device)
# Compute score
score = self.model(speech.unsqueeze(0), self.sample_rate)
scores.append(score.item())
return scores
def score_dir(self, dir_path: str, extension: str) -> float:
"""
Computes the average UTMOS score for all files in a directory.
Args:
dir_path (str): Path to the directory containing audio files.
Returns:
float: Average UTMOS score for the directory.
"""
logging.info(f"Calculating UTMOS score for {dir_path}")
# Get list of wav files
wav_files = [
os.path.join(dir_path, f)
for f in os.listdir(dir_path)
if f.lower().endswith(extension)
]
if not wav_files:
raise ValueError(f"No audio files found in {dir_path}")
# Compute scores
scores = self.score_files(wav_files)
return float(np.mean(scores))
if __name__ == "__main__":
torch.set_num_threads(1)
torch.set_num_interop_threads(1)
formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"
logging.basicConfig(format=formatter, level=logging.INFO, force=True)
parser = get_parser()
args = parser.parse_args()
# Validate input path
if not os.path.isdir(args.wav_path):
logging.error(f"Invalid directory: {args.wav_path}")
exit(1)
# Initialize evaluator
model_path = os.path.join(args.model_dir, "mos/utmos22_strong_step7459_v1.pt")
if not os.path.exists(model_path):
logging.error(
"Please download evaluation models from "
"https://huggingface.co/k2-fsa/TTS_eval_models"
" and pass this dir with --model-dir"
)
exit(1)
utmos_evaluator = UTMOSScore(model_path)
# Compute UTMOS score
score = utmos_evaluator.score_dir(args.wav_path, args.extension)
print("-" * 50)
logging.info(f"UTMOS score: {score:.2f}")
print("-" * 50)
|