scripts/run_phase1_on_folder.py

"""Run Phase 1 on every WAV in a folder, writing per-segment JSON.

Bypasses the annotation-tool DB — produces standalone JSON output for one-off
tests. Reuses the same `annotate.phase1` modules so the analysis is identical
to what the annotation tool produces.

Usage:
    python scripts/run_phase1_on_folder.py SEGMENTS_DIR OUT_DIR LANGUAGE

Example:
    python scripts/run_phase1_on_folder.py \\
        data/hindi_emphasis_test/segments \\
        data/hindi_emphasis_test/phase1 \\
        hi
"""

from __future__ import annotations

import json
import logging
import sys
from pathlib import Path

# Make annotate/ importable when running this script from anywhere
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import time

from annotate.phase1.config import PIPELINE_VERSION
from annotate.phase1.emphasis import score_emphasis
from annotate.phase1.emotion import classify_emotion, offload_emotion_model
from annotate.phase1.prosody import extract_per_word_prosody
from annotate.phase1.transcribe import transcribe

log = logging.getLogger(__name__)


def process_one(wav_path: Path, language: str) -> dict:
    """Run the full Phase 1 pipeline on one WAV; return a JSON-serializable dict."""
    t0 = time.time()

    # 1. Transcribe (WhisperX + forced alignment).
    asr = transcribe(wav_path, language=language)

    # 2. Prosody per word (parselmouth F0 + energy).
    windows = [(w.start, w.end) for w in asr.words]
    prosody = extract_per_word_prosody(wav_path, windows)
    f0_peaks = [p[0] for p in prosody]
    energy_peaks = [p[1] for p in prosody]

    # 3. Emphasis flags.
    duration = windows[-1][1] if windows else 0.0
    flags = score_emphasis(windows, f0_peaks, energy_peaks, duration)

    # 4. Emotion (segment-level via emotion2vec+).
    try:
        emo = classify_emotion(wav_path)
        emotion_obj = {
            "label": emo.label,
            "confidence": emo.confidence,
            "scores": emo.scores,
        }
    except Exception as e:
        emotion_obj = {"error": f"{type(e).__name__}: {e}"}

    elapsed = round(time.time() - t0, 2)

    return {
        "segment": wav_path.name,
        "language": asr.language,
        "duration_seconds": round(duration, 3),
        "text": asr.text,
        "n_words": len(asr.words),
        "n_emphasized": int(sum(flags)),
        "words": [
            {
                "idx": i,
                "text": w.text,
                "start": round(w.start, 3),
                "end": round(w.end, 3),
                "emphasis": bool(flags[i]),
                "f0_peak": round(f0_peaks[i], 1) if f0_peaks[i] is not None else None,
                "energy_peak": round(energy_peaks[i], 1) if energy_peaks[i] is not None else None,
            }
            for i, w in enumerate(asr.words)
        ],
        "emotion": emotion_obj,
        "phase1_version": PIPELINE_VERSION,
        "elapsed_seconds": elapsed,
    }


def main(seg_dir: Path, out_dir: Path, language: str) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    wavs = sorted(p for p in seg_dir.glob("*.wav") if not p.name.startswith("_"))
    if not wavs:
        print(f"No WAVs found in {seg_dir}")
        return

    print(f"Processing {len(wavs)} segments (language={language})...")
    for wav in wavs:
        try:
            result = process_one(wav, language=language)
            out_path = out_dir / (wav.stem + ".json")
            out_path.write_text(json.dumps(result, indent=2, ensure_ascii=False), encoding="utf-8")
            print(
                f"  {wav.name}: {result['n_words']:3d} words, "
                f"{result['n_emphasized']} emphasized, "
                f"emotion={result['emotion'].get('label', 'err')}, "
                f"{result['elapsed_seconds']}s"
            )
        except Exception as e:
            log.exception("phase1 failed for %s", wav)
            err_path = out_dir / (wav.stem + ".error.json")
            err_path.write_text(json.dumps({"error": str(e)}, indent=2), encoding="utf-8")
            print(f"  {wav.name}: FAILED -- {e}")

    # Free emotion model GPU memory at the end
    try:
        offload_emotion_model()
    except Exception:
        pass

    print(f"Done. Outputs in {out_dir}")


if __name__ == "__main__":
    if len(sys.argv) != 4:
        print("Usage: run_phase1_on_folder.py SEGMENTS_DIR OUT_DIR LANGUAGE")
        sys.exit(1)

    logging.basicConfig(level=logging.WARNING, format="%(levelname)s %(name)s: %(message)s")
    for n in ("transformers", "pyannote", "pytorch_lightning", "huggingface_hub", "modelscope", "funasr", "root"):
        logging.getLogger(n).setLevel(logging.ERROR)

    main(Path(sys.argv[1]), Path(sys.argv[2]), sys.argv[3])