omnivoice/eval/mos/utmos.py

#!/usr/bin/env python3
# Copyright    2026  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 multiprocessing as mp
import os
import sys
import traceback
import warnings
from concurrent.futures import ProcessPoolExecutor, as_completed

import numpy as np
import torch
from tqdm import tqdm

from omnivoice.eval.models.utmos import UTMOS22Strong
from omnivoice.eval.utils import load_eval_waveform
from omnivoice.utils.data_utils import read_test_list

warnings.filterwarnings("ignore")

# Global variables for workers
worker_model = None
worker_device = None
worker_sr = 16000


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(
        "--test-list",
        type=str,
        required=True,
        help="Path to the JSONL test list. Each line is a JSON object "
        "with fields: id, text, ref_audio, ref_text, language_id, language_name.",
    )
    parser.add_argument(
        "--model-dir",
        type=str,
        required=True,
        help="Local path of our evaluation 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",
    )
    parser.add_argument(
        "--decode-path",
        type=str,
        default=None,
        help="Path to the output file where UTMOS information will be saved. "
        "If not provided, results are only printed to console.",
    )
    parser.add_argument(
        "--nj-per-gpu",
        type=int,
        default=1,
        help="Number of worker processes to spawn per GPU.",
    )
    return parser


def get_device(rank: int = 0) -> torch.device:
    assert torch.cuda.is_available(), "CUDA is required but not available."
    device = torch.device(f"cuda:{rank}")
    torch.cuda.set_device(rank)
    return device


def worker_init(
    rank_queue,
    model_path,
):
    """Initialize worker process with model and device."""
    global worker_model, worker_device, worker_sr

    # Limit CPU threads per worker
    torch.set_num_threads(2)

    formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] [Worker %(process)d] %(message)s"
    logging.basicConfig(format=formatter, level=logging.INFO, force=True)

    rank = rank_queue.get() if rank_queue else -1

    worker_device = get_device(rank)
    worker_sr = 16000

    logging.debug(f"Initializing UTMOS worker on {worker_device}")

    # Initialize Model
    worker_model = UTMOS22Strong()
    try:
        # Load weights to CPU first, then move to device
        state_dict = torch.load(model_path, map_location="cpu")
        worker_model.load_state_dict(state_dict)
    except Exception as e:
        logging.error(f"Failed to load model from {model_path}: {e}")
        raise

    worker_model.to(worker_device)
    worker_model.eval()


@torch.no_grad()
def run_utmos_worker(file_idx, wav_path, language_name):
    """Worker function to process a single audio file."""
    try:
        if not os.path.exists(wav_path):
            return file_idx, wav_path, language_name, f"File not found: {wav_path}", "error"

        # Load and preprocess waveform
        speech = load_eval_waveform(wav_path, worker_sr, device=worker_device)

        # Compute score
        # UTMOS expects input shape (Batch, Time)
        score = worker_model(speech.unsqueeze(0), worker_sr)

        return file_idx, wav_path, language_name, score.item(), "success"

    except Exception as e:
        error_detail = (
            f"Error processing {wav_path}: {str(e)}\n"
            f"Traceback:\n{traceback.format_exc()}"
        )
        return file_idx, wav_path, language_name, error_detail, "error"


def main():
    parser = get_parser()
    args = parser.parse_args()

    # Main process thread setting
    torch.set_num_threads(2)

    mp.set_start_method("spawn", force=True)

    formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"
    logging.basicConfig(format=formatter, level=logging.INFO, force=True)

    # Validate inputs
    if not os.path.isdir(args.wav_path):
        logging.error(f"Invalid directory: {args.wav_path}")
        sys.exit(1)

    model_path = os.path.join(args.model_dir, "mos/utmos22_strong_step7459_v1.pt")
    if not os.path.exists(model_path):
        logging.error(f"Model file not found at {model_path}")
        sys.exit(1)

    # Scan directory for files
    logging.info(f"Calculating UTMOS for {args.wav_path}")

    wav_files = []
    try:
        samples = read_test_list(args.test_list)
        for s in samples:
            language_name = s.get("language_name") or "unknown"
            eval_wav_path = os.path.join(args.wav_path, f"{s['id']}.{args.extension}")
            wav_files.append((eval_wav_path, language_name))
    except Exception as e:
        raise ValueError(f"Error reading test list {args.test_list}: {e}")

    # Setup Parallel Processing
    num_gpus = torch.cuda.device_count()
    assert num_gpus > 0, "No GPU found. GPU is required."
    total_procs = num_gpus * args.nj_per_gpu

    logging.info(
        f"Starting evaluation with {total_procs} processes on {num_gpus} GPUs."
    )

    manager = mp.Manager()
    rank_queue = manager.Queue()

    for rank in list(range(num_gpus)) * args.nj_per_gpu:
        rank_queue.put(rank)

    scores = []

    fout = None
    if args.decode_path:
        os.makedirs(os.path.dirname(args.decode_path), exist_ok=True)
        fout = open(args.decode_path, "w", encoding="utf8")
        logging.info(f"Saving detailed UTMOS results to: {args.decode_path}")
        fout.write("Name\tUTMOS\n")

    try:
        with ProcessPoolExecutor(
            max_workers=total_procs,
            initializer=worker_init,
            initargs=(
                rank_queue,
                model_path,
            ),
        ) as executor:
            futures = []
            for i, (wav_path, language_name) in enumerate(wav_files):
                futures.append(
                    executor.submit(run_utmos_worker, i, wav_path, language_name)
                )

            pbar = tqdm(
                as_completed(futures), total=len(wav_files), desc="Evaluating UTMOS"
            )
            lang_stats = {}
            for future in pbar:
                idx, path, language_name, result, status = future.result()
                if status == "success":
                    if language_name not in lang_stats:
                        lang_stats[language_name] = []
                    lang_stats[language_name].append(result)
                    scores.append(result)
                    if fout:
                        if language_name == "unknown":
                            fout.write(f"{os.path.basename(path)}\t{result:.2f}\n")
                        else:
                            fout.write(
                                f"{language_name}\t{os.path.basename(path)}\t{result:.2f}\n"
                            )
                else:
                    pbar.write(f"!!! FAILED [File {idx}]: {path} | {result}")

    except (Exception, KeyboardInterrupt) as e:
        logging.critical(
            f"An unrecoverable error occurred: {e}. Terminating all processes."
        )
        detailed_error_info = traceback.format_exc()
        logging.error(f"--- DETAILED TRACEBACK ---\n{detailed_error_info}")
        sys.exit(1)

    print("-" * 50)

    if len(lang_stats) > 1:
        lang_scores = []
        for lang in sorted(lang_stats.keys()):
            l_scores = lang_stats[lang]
            l_avg = np.mean(l_scores)
            lang_scores.append(l_scores)
            l_count = len(l_scores)
            logging.info(f"[{lang}] UTMOS score: {l_avg:.3f} ({l_count} samples)")
            if fout:
                fout.write(f"[{lang}] UTMOS: {l_avg:.3f} ({l_count} samples)\n")
        logging.info(
            f"Macro-average UTMOS over {len(lang_stats)} languages: "
            f"{np.mean([np.mean(ls) for ls in lang_scores]):.3f}"
        )
        if fout:
            fout.write(
                f"\nMacro-average UTMOS over {len(lang_stats)} languages: "
                f"{np.mean([np.mean(ls) for ls in lang_scores]):.3f}\n"
            )

    if scores:
        avg_score = np.mean(scores)
        logging.info(f"Processed {len(scores)}/{len(wav_files)} files.")
        logging.info(f"UTMOS score: {avg_score:.2f}")
        if fout:
            fout.write(f"\nAverage UTMOS: {avg_score:.2f}\n")
    else:
        logging.error("No valid scores computed.")
    print("-" * 50)

    if fout:
        fout.close()


if __name__ == "__main__":
    main()