| 1 | import argparse |
| 2 | import json |
| 3 | from pathlib import Path |
| 4 | |
| 5 | import librosa |
| 6 | import torch |
| 7 | from tqdm import tqdm |
| 8 | |
| 9 | |
| 10 | def main(): |
| 11 | parser = argparse.ArgumentParser(description="UTMOS Evaluation") |
| 12 | parser.add_argument("--audio_dir", type=str, required=True, help="Audio file path.") |
| 13 | parser.add_argument("--ext", type=str, default="wav", help="Audio extension.") |
| 14 | args = parser.parse_args() |
| 15 | |
| 16 | device = "cuda" if torch.cuda.is_available() else "xpu" if torch.xpu.is_available() else "cpu" |
| 17 | |
| 18 | predictor = torch.hub.load("tarepan/SpeechMOS:v1.2.0", "utmos22_strong", trust_repo=True) |
| 19 | predictor = predictor.to(device) |
| 20 | |
| 21 | audio_paths = list(Path(args.audio_dir).rglob(f"*.{args.ext}")) |
| 22 | utmos_score = 0 |
| 23 | |
| 24 | utmos_result_path = Path(args.audio_dir) / "_utmos_results.jsonl" |
| 25 | with open(utmos_result_path, "w", encoding="utf-8") as f: |
| 26 | for audio_path in tqdm(audio_paths, desc="Processing"): |
| 27 | wav, sr = librosa.load(audio_path, sr=None, mono=True) |
| 28 | wav_tensor = torch.from_numpy(wav).to(device).unsqueeze(0) |
| 29 | score = predictor(wav_tensor, sr) |
| 30 | line = {} |
| 31 | line["wav"], line["utmos"] = str(audio_path.stem), score.item() |
| 32 | utmos_score += score.item() |
| 33 | f.write(json.dumps(line, ensure_ascii=False) + "\n") |
| 34 | avg_score = utmos_score / len(audio_paths) if len(audio_paths) > 0 else 0 |
| 35 | f.write(f"\nUTMOS: {avg_score:.4f}\n") |
| 36 | |
| 37 | print(f"UTMOS: {avg_score:.4f}") |
| 38 | print(f"UTMOS results saved to {utmos_result_path}") |
| 39 | |
| 40 | |
| 41 | if __name__ == "__main__": |
| 42 | main() |
| 43 |