inference_example.py

#!/usr/bin/env python3
"""Inference example for Speech Artifact Detectors.

Demonstrates loading models from HuggingFace Hub and scoring audio files.

Usage:
    # Score a single file with all 10 detectors
    python inference_example.py audio.wav

    # Score a directory of files
    python inference_example.py /path/to/audio/ --ext wav

    # Only TTS/vocoder detectors
    python inference_example.py audio.wav --category tts_artifact

    # Only augmentation detectors on CPU
    python inference_example.py audio.wav --category augmentation --device cpu

    # Custom threshold for flagging artifacts
    python inference_example.py audio.wav --threshold 0.3

Requirements:
    pip install torch torchaudio soundfile numpy huggingface_hub
"""

import argparse
import glob
import os
import sys
import time

import numpy as np
import torch


def main():
    parser = argparse.ArgumentParser(
        description="Score audio files for speech artifacts",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("input", help="Audio file or directory to score")
    parser.add_argument("--ext", default="wav",
                        help="File extension when input is a directory (default: wav)")
    parser.add_argument("--device", default=None,
                        help="Device: cuda, cpu, or auto-detect (default)")
    parser.add_argument("--category", choices=["tts_artifact", "augmentation"],
                        help="Load only models from this category")
    parser.add_argument("--threshold", type=float, default=0.5,
                        help="Score threshold for flagging artifacts (default: 0.5)")
    parser.add_argument("--batch-size", type=int, default=16,
                        help="Batch size for inference (default: 16)")
    args = parser.parse_args()

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

    # ── Load models from HuggingFace Hub ─────────────────────────────────
    from speech_artifact_detector import load_from_hub

    print(f"Loading models on {device}...")
    categories = [args.category] if args.category else None
    t0 = time.time()
    models = load_from_hub(device=device, categories=categories)
    print(f"  Loaded {len(models)} models in {time.time() - t0:.1f}s: "
          f"{', '.join(models.keys())}\n")

    # ── Collect input files ──────────────────────────────────────────────
    if os.path.isdir(args.input):
        files = sorted(glob.glob(os.path.join(args.input, f"*.{args.ext}")))
        if not files:
            print(f"No .{args.ext} files found in {args.input}")
            sys.exit(1)
    elif os.path.isfile(args.input):
        files = [args.input]
    else:
        print(f"File not found: {args.input}")
        sys.exit(1)

    print(f"Scoring {len(files)} file(s)...\n")

    # ── Score each file ──────────────────────────────────────────────────
    from speech_artifact_detector import score_file_all

    model_names = list(models.keys())
    all_results = []

    for fpath in files:
        fname = os.path.basename(fpath)
        scores = score_file_all(models, fpath, device)
        mean_score = np.mean(list(scores.values()))
        verdict = "ARTIFACT" if mean_score >= args.threshold else "clean"
        all_results.append({"file": fname, "scores": scores,
                            "mean": mean_score, "verdict": verdict})

        # Print per-file results
        print(f"{'=' * 60}")
        print(f"  {fname}  →  {verdict} (mean={mean_score:.4f})")
        print(f"{'=' * 60}")
        for name in model_names:
            s = scores[name]
            bar = "#" * int(s * 30) + "." * (30 - int(s * 30))
            flag = " <<<" if s >= args.threshold else ""
            print(f"  {name:>25s}: {s:.4f}  [{bar}]{flag}")
        print()

    # ── Summary ──────────────────────────────────────────────────────────
    if len(files) > 1:
        n_artifact = sum(1 for r in all_results if r["verdict"] == "ARTIFACT")
        print(f"\nSummary: {n_artifact}/{len(files)} files flagged as ARTIFACT "
              f"(threshold={args.threshold})")
        print(f"\nPer-detector averages across all files:")
        for name in model_names:
            avg = np.mean([r["scores"][name] for r in all_results])
            print(f"  {name:>25s}: {avg:.4f}")


if __name__ == "__main__":
    main()