tools/maps/maps_align_json.py

#!/usr/bin/env python3
"""JSON-manifest MAPS alignment for inpaint_pipeline.py --maps.

Reads segment transcript JSON (``02_transcripts.json``), runs MAPS in the
active PhonePaint interpreter (no conda hop, no TensorFlow), and writes
per-segment alignments with PhonePaint ``canonical`` labels.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

from textgrid import IntervalTier, TextGrid

MAPS_ROOT = Path(__file__).resolve().parent
REPO_ROOT = MAPS_ROOT.parents[1]
sys.path.insert(0, str(MAPS_ROOT))
sys.path.insert(0, str(REPO_ROOT))

from phone_labels import canonical_phone_label  # noqa: E402
from maps_torch.align_runtime import (  # noqa: E402
    align_wav,
    load_models,
    resolve_device,
)
from maps_torch.alignment import load_dictionary  # noqa: E402

DEFAULT_MAPS_CKPT = MAPS_ROOT / "torch_models" / "timbuck_eng.pt"
DEFAULT_MAPS_ENSEMBLE = MAPS_ROOT / "torch_models" / "ensemble_model"
_BUNDLED_CMUDICT = MAPS_ROOT / "nltk_data" / "corpora" / "cmudict" / "cmudict"
_NLTK_CMUDICT = Path.home() / "nltk_data" / "corpora" / "cmudict" / "cmudict"


def _default_cmudict() -> Path:
    env_value = os.environ.get("PHONEPAINT_CMUDICT")
    if env_value:
        env_path = Path(env_value).expanduser()
        if env_path.is_file():
            return env_path
        # A stale env var (e.g. an unmounted drive) must not break a good checkout.
        print(
            f"PHONEPAINT_CMUDICT={env_value} is not a file — ignoring it.",
            file=sys.stderr,
        )
    for candidate in (_BUNDLED_CMUDICT, _NLTK_CMUDICT):
        if candidate.is_file():
            return candidate
    return _BUNDLED_CMUDICT


DEFAULT_MAPS_DICT = _default_cmudict()


def _pick_tiers(tg: TextGrid) -> tuple[IntervalTier | None, IntervalTier | None]:
    phones_tier = None
    words_tier = None
    for tier in tg.tiers:
        if not isinstance(tier, IntervalTier):
            continue
        lname = (tier.name or "").lower()
        if "phone" in lname or "segment" in lname:
            phones_tier = tier
        if "word" in lname:
            words_tier = tier
    return phones_tier, words_tier


def _tier_to_dict(tier: IntervalTier | None, *, phones: bool) -> dict[str, dict]:
    out: dict[str, dict] = {}
    if tier is None:
        return out
    for interval in tier:
        text = (interval.mark or "").strip()
        if not text:
            continue
        start = round(float(interval.minTime), 3)
        end = round(float(interval.maxTime), 3)
        if end <= start:
            continue
        item: dict = {"xmin": start, "xmax": end, "text": text}
        if phones:
            item["canonical"] = canonical_phone_label(text)
            item["inventory"] = "maps-timit61"
        out[str(len(out))] = item
    return out


def _textgrid_to_alignment(tg: TextGrid) -> dict:
    phones_tier, words_tier = _pick_tiers(tg)
    phones = _tier_to_dict(phones_tier, phones=True)
    if not phones:
        raise RuntimeError("MAPS produced no phone intervals")
    return {
        "words": _tier_to_dict(words_tier, phones=False),
        "phones": phones,
        "alignment": {
            "backend": "maps",
            "inventory": "maps-timit61",
        },
    }


def _load_transcripts(path: Path) -> dict[str, str]:
    raw = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(raw, dict):
        raise TypeError(f"Expected transcript object in {path}")
    texts: dict[str, str] = {}
    for name, entry in raw.items():
        if isinstance(entry, dict):
            text = str(entry.get("text") or "").strip()
        else:
            text = str(entry or "").strip()
        if not text:
            raise ValueError(f"Missing transcript text for {name} in {path}")
        texts[str(name)] = text
    return texts


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--audio_dir", type=Path, required=True)
    ap.add_argument("--input_json", type=Path, required=True)
    ap.add_argument("--output_json", type=Path, required=True)
    ap.add_argument(
        "--model",
        type=Path,
        default=None,
        help="Single .pt checkpoint or ensemble directory (default: timbuck_eng.pt)",
    )
    ap.add_argument(
        "--dict",
        dest="dictionary",
        type=Path,
        default=None,
        help=f"CMUdict path (default: PHONEPAINT_CMUDICT or {DEFAULT_MAPS_DICT})",
    )
    ap.add_argument(
        "--device",
        default="auto",
        choices=("auto", "cpu", "cuda"),
        help="Torch device for MAPS acoustic model (default: auto)",
    )
    ap.add_argument(
        "--no-sil",
        action="store_true",
        help="Disable silence padding at transcript ends",
    )
    ap.add_argument(
        "--no-interp",
        action="store_true",
        help="Disable boundary interpolation",
    )
    ap.add_argument(
        "--check-variants",
        action="store_true",
        help="Evaluate pronunciation variants (slow)",
    )
    ap.add_argument(
        "--ensemble",
        action="store_true",
        help="Use tools/maps/torch_models/ensemble_model/ (10 checkpoints, median boundaries)",
    )
    args = ap.parse_args()

    model_path = args.model
    if model_path is None:
        model_path = DEFAULT_MAPS_ENSEMBLE if args.ensemble else DEFAULT_MAPS_CKPT
    model_path = model_path.expanduser().resolve()
    if not model_path.exists():
        raise SystemExit(
            f"MAPS checkpoint not found: {model_path}\n"
            "Run tools/maps/sync_checkpoints.sh or copy torch_models/ from ~/MAPS."
        )

    dict_path = (args.dictionary or DEFAULT_MAPS_DICT).expanduser().resolve()
    if not dict_path.is_file():
        raise SystemExit(
            f"CMUdict not found: {dict_path}\n"
            "Set PHONEPAINT_CMUDICT or pass --dict."
        )

    device = resolve_device(args.device)
    print(f"MAPS device: {device}", flush=True)
    models = load_models(model_path, device)
    print(f"Loaded {len(models)} MAPS checkpoint(s) from {model_path}", flush=True)
    word2phone = load_dictionary(dict_path)

    audio_dir = args.audio_dir.expanduser().resolve()
    transcripts = _load_transcripts(args.input_json.expanduser().resolve())

    aligned: dict[str, dict] = {}
    for name, text in transcripts.items():
        wav_path = audio_dir / name
        if not wav_path.is_file():
            raise FileNotFoundError(f"Segment wav not found: {wav_path}")
        print(f"MAPS align: {name}", flush=True)
        tg = align_wav(
            wav_path,
            text,
            models=models,
            word2phone=word2phone,
            device=device,
            add_sil=not args.no_sil,
            use_interp=not args.no_interp,
            check_variants=args.check_variants,
        )
        aligned[name] = _textgrid_to_alignment(tg)

    args.output_json.parent.mkdir(parents=True, exist_ok=True)
    args.output_json.write_text(
        json.dumps(aligned, indent=2, ensure_ascii=False),
        encoding="utf-8",
    )
    print(f"Wrote {args.output_json} ({len(aligned)} segment(s))", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())