evaluate.py

"""
Evaluation pipeline for PromptTTS.

Computes objective metrics on the test set:
- Mel Cepstral Distortion (MCD)
- F0 RMSE
- Style classification accuracy

Usage:
    python evaluate.py \
        --checkpoint checkpoints/best_model.pt \
        --test_filelist data/processed/filelists/test.txt \
        --features_dir data/processed/features
"""

import argparse
import json
import os

import numpy as np
import torch
from tqdm import tqdm

from models.prompttts import PromptTTS
from data.dataset import create_dataloader
from utils.audio import AudioProcessor
from utils.metrics import mel_cepstral_distortion, f0_rmse, style_classification_accuracy


def evaluate(
    checkpoint_path: str,
    test_filelist: str,
    features_dir: str,
    device: torch.device,
    output_dir: str = "eval_results",
):
    os.makedirs(output_dir, exist_ok=True)

    # Load model
    print("Loading model...")
    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
    config = checkpoint["config"]

    model = PromptTTS(
        style_config=config["model"]["style_encoder"],
        content_config=config["model"]["content_encoder"],
        decoder_config=config["model"]["decoder"],
    )
    model.load_state_dict(checkpoint["model"])
    model.to(device)
    model.eval()

    # DataLoader
    test_loader = create_dataloader(
        filelist_path=test_filelist,
        features_dir=features_dir,
        batch_size=1,
        shuffle=False,
    )

    # Evaluation
    all_mcd = []
    all_f0_rmse = []
    all_style_logits = {k: [] for k in ["gender", "pitch", "speed", "volume", "emotion"]}
    all_style_labels = {k: [] for k in ["gender", "pitch", "speed", "volume", "emotion"]}

    audio_processor = AudioProcessor()

    print("Running evaluation...")
    for batch in tqdm(test_loader):
        batch = {
            k: v.to(device) if isinstance(v, torch.Tensor) else v
            for k, v in batch.items()
        }
        if "style_labels" in batch:
            batch["style_labels"] = {
                k: v.to(device) for k, v in batch["style_labels"].items()
            }

        with torch.no_grad():
            outputs = model(
                style_input_ids=batch["style_input_ids"],
                style_attention_mask=batch["style_attention_mask"],
                phonemes=batch["phonemes"],
                phoneme_mask=batch["phoneme_mask"],
                pitch_target=batch["pitch"],
                energy_target=batch["energy"],
                mel_mask=batch["mel_mask"],
            )

        # Mel Cepstral Distortion
        pred_mel = outputs["mel_output"].squeeze(0).cpu().numpy()
        ref_mel = batch["mel"].squeeze(0).cpu().numpy()
        mcd = mel_cepstral_distortion(pred_mel, ref_mel)
        all_mcd.append(mcd)

        # F0 RMSE
        pred_pitch = outputs["pitch_pred"].squeeze(0).cpu().numpy()
        ref_pitch = batch["pitch"].squeeze(0).cpu().numpy()
        f0_err = f0_rmse(pred_pitch, ref_pitch)
        all_f0_rmse.append(f0_err)

        # Style classification
        for key in all_style_logits:
            if key in outputs["style_logits"]:
                all_style_logits[key].append(outputs["style_logits"][key].cpu())
                all_style_labels[key].append(batch["style_labels"][key].cpu())

    # Aggregate metrics
    results = {
        "mcd_mean": float(np.mean(all_mcd)),
        "mcd_std": float(np.std(all_mcd)),
        "f0_rmse_mean": float(np.mean(all_f0_rmse)),
        "f0_rmse_std": float(np.std(all_f0_rmse)),
        "num_samples": len(all_mcd),
    }

    # Style classification accuracy
    if all_style_logits["gender"]:
        style_logits_cat = {
            k: torch.cat(v, dim=0) for k, v in all_style_logits.items() if v
        }
        style_labels_cat = {
            k: torch.cat(v, dim=0) for k, v in all_style_labels.items() if v
        }
        style_acc = style_classification_accuracy(style_logits_cat, style_labels_cat)
        results["style_accuracy"] = style_acc

    # Print results
    print("\n" + "=" * 50)
    print("EVALUATION RESULTS")
    print("=" * 50)
    print(f"  MCD:          {results['mcd_mean']:.3f} +/- {results['mcd_std']:.3f} dB")
    print(f"  F0 RMSE:      {results['f0_rmse_mean']:.3f} +/- {results['f0_rmse_std']:.3f} Hz")
    if "style_accuracy" in results:
        print(f"  Style Accuracy:")
        for k, v in results["style_accuracy"].items():
            print(f"    {k:>10s}: {v:.3f}")
    print(f"  Num samples:  {results['num_samples']}")
    print("=" * 50)

    # Save results
    results_path = os.path.join(output_dir, "eval_results.json")
    with open(results_path, "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nResults saved to {results_path}")


def main():
    parser = argparse.ArgumentParser(description="Evaluate PromptTTS")
    parser.add_argument("--checkpoint", type=str, required=True)
    parser.add_argument("--test_filelist", type=str, default="data/processed/filelists/test.txt")
    parser.add_argument("--features_dir", type=str, default="data/processed/features")
    parser.add_argument("--output_dir", type=str, default="eval_results")
    parser.add_argument("--device", type=str, default="auto")
    args = parser.parse_args()

    if args.device == "auto":
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    else:
        device = torch.device(args.device)

    evaluate(
        checkpoint_path=args.checkpoint,
        test_filelist=args.test_filelist,
        features_dir=args.features_dir,
        device=device,
        output_dir=args.output_dir,
    )


if __name__ == "__main__":
    main()