| import argparse |
| import json |
| from pathlib import Path |
|
|
| import librosa |
| import torch |
| from tqdm import tqdm |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="UTMOS Evaluation") |
| parser.add_argument("--audio_dir", type=str, required=True, help="Audio file path.") |
| parser.add_argument("--ext", type=str, default="wav", help="Audio extension.") |
| args = parser.parse_args() |
|
|
| device = "cuda" if torch.cuda.is_available() else "xpu" if torch.xpu.is_available() else "cpu" |
|
|
| predictor = torch.hub.load("tarepan/SpeechMOS:v1.2.0", "utmos22_strong", trust_repo=True) |
| predictor = predictor.to(device) |
|
|
| audio_paths = list(Path(args.audio_dir).rglob(f"*.{args.ext}")) |
| utmos_score = 0 |
|
|
| utmos_result_path = Path(args.audio_dir) / "_utmos_results.jsonl" |
| with open(utmos_result_path, "w", encoding="utf-8") as f: |
| for audio_path in tqdm(audio_paths, desc="Processing"): |
| wav, sr = librosa.load(audio_path, sr=None, mono=True) |
| wav_tensor = torch.from_numpy(wav).to(device).unsqueeze(0) |
| score = predictor(wav_tensor, sr) |
| line = {} |
| line["wav"], line["utmos"] = str(audio_path.stem), score.item() |
| utmos_score += score.item() |
| f.write(json.dumps(line, ensure_ascii=False) + "\n") |
| avg_score = utmos_score / len(audio_paths) if len(audio_paths) > 0 else 0 |
| f.write(f"\nUTMOS: {avg_score:.4f}\n") |
|
|
| print(f"UTMOS: {avg_score:.4f}") |
| print(f"UTMOS results saved to {utmos_result_path}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
