Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """Score generated speech with UTMOS. | |
| This is an auxiliary non-intrusive objective proxy. It uses the public | |
| ``tarepan/SpeechMOS`` torch.hub model, writes per-utterance JSONL, and keeps the | |
| metric separate from human MOS/CMOS evidence. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import statistics | |
| import sys | |
| from pathlib import Path | |
| from typing import Any | |
| def parse_args() -> argparse.Namespace: | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--audio_dir", type=Path, required=True) | |
| parser.add_argument("--out_json", type=Path, required=True) | |
| parser.add_argument("--ext", default="wav") | |
| parser.add_argument("--max_items", type=int, default=0) | |
| parser.add_argument("--device", default="") | |
| return parser.parse_args() | |
| def _mean(values: list[float]) -> float | None: | |
| return statistics.mean(values) if values else None | |
| def _std(values: list[float]) -> float | None: | |
| if not values: | |
| return None | |
| if len(values) == 1: | |
| return 0.0 | |
| return statistics.stdev(values) | |
| def _wav_paths(audio_dir: Path, ext: str, max_items: int) -> list[Path]: | |
| wavs = sorted(audio_dir.rglob(f"*.{ext}")) | |
| if max_items > 0: | |
| wavs = wavs[:max_items] | |
| return wavs | |
| def main() -> int: | |
| args = parse_args() | |
| wavs = _wav_paths(args.audio_dir, args.ext, args.max_items) | |
| print(f"[utmos] {len(wavs)} wavs from {args.audio_dir}", file=sys.stderr) | |
| try: | |
| import librosa | |
| import torch | |
| except Exception as exc: # pragma: no cover - optional dependency path | |
| raise SystemExit("UTMOS dependencies are missing: install torch and librosa.") from exc | |
| device = args.device | |
| if not device: | |
| 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).eval() | |
| rows: list[dict[str, Any]] = [] | |
| n_skip = 0 | |
| with torch.no_grad(): | |
| for idx, wav in enumerate(wavs, start=1): | |
| try: | |
| audio, sr = librosa.load(str(wav), sr=None, mono=True) | |
| wav_tensor = torch.from_numpy(audio).to(device).unsqueeze(0) | |
| score = predictor(wav_tensor, sr) | |
| score_float = float(score.item() if hasattr(score, "item") else score) | |
| rows.append( | |
| { | |
| "wav": str(wav), | |
| "item_id": wav.stem, | |
| "utmos": score_float, | |
| } | |
| ) | |
| except Exception as exc: # pragma: no cover - data-dependent path | |
| n_skip += 1 | |
| print(f"[utmos-skip] {wav}: {exc}", file=sys.stderr) | |
| if idx % 25 == 0 or idx == len(wavs): | |
| print(f"[utmos] {idx}/{len(wavs)}", file=sys.stderr) | |
| values = [float(row["utmos"]) for row in rows] | |
| out = { | |
| "audio_dir": str(args.audio_dir), | |
| "device": device, | |
| "max_items": args.max_items, | |
| "metric": "utmos22_strong", | |
| "model": "tarepan/SpeechMOS:v1.2.0 utmos22_strong", | |
| "n_items": len(wavs), | |
| "n_scored": len(rows), | |
| "n_skipped": n_skip, | |
| "utmos_mean": _mean(values), | |
| "utmos_std": _std(values), | |
| } | |
| args.out_json.parent.mkdir(parents=True, exist_ok=True) | |
| args.out_json.write_text(json.dumps(out, indent=2, sort_keys=True) + "\n", encoding="utf-8") | |
| per_utt_path = args.out_json.parent / f"{args.out_json.stem}_per_utt_utmos_objective.jsonl" | |
| per_utt_path.write_text( | |
| "\n".join(json.dumps(row, sort_keys=True) for row in rows) + "\n", | |
| encoding="utf-8", | |
| ) | |
| print(json.dumps(out, indent=2, sort_keys=True)) | |
| return 0 | |
| if __name__ == "__main__": | |
| raise SystemExit(main()) | |