reference/diarize.py

#!/usr/bin/env python
# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""Anuvani-Diarization (Indic model, release 1.0) — audio in, RTTM out.

  python diarize.py --out OUT_DIR audio1.wav [audio2.flac ...]
  python diarize.py --out OUT_DIR --in-dir AUDIO_DIR

Utterance-style segments matching IndicDiarBench's labelling convention (cannot-link VBx clustering and a learned
decoder; see README.md).
Determinism: FP32, TF32 off, deterministic kernels, seed 3407 per file. Byte-identical outputs to the published
reference RTTMs were verified with the defaults (--batch-size 8, --threads 2); see tools/verify_golden.py.
"""
import argparse
import hashlib
import json
import platform
import sys
import time
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))

AUDIO_EXT = ("*.wav", "*.flac", "*.mp3", "*.m4a", "*.ogg")


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("audio", nargs="*", type=Path)
    p.add_argument("--in-dir", type=Path)
    p.add_argument("--out", type=Path, required=True)
    p.add_argument("--bundle", type=Path, default=HERE / "bundle")
    p.add_argument("--device", default="cuda:0", help="cuda:N or cpu")
    p.add_argument("--batch-size", type=int, default=8)
    p.add_argument("--threads", type=int, default=2)
    a = p.parse_args()
    files = list(a.audio)
    if a.in_dir:
        for ext in AUDIO_EXT:
            files.extend(sorted(a.in_dir.glob(ext)))
    files = [f for f in files if f.is_file()]
    if not files:
        raise SystemExit("no audio files given")
    from anuvani_diar_ref.frontend import load_bundle, load_audio, analyse
    from anuvani_diar_ref.modes import indic_mode, write_rttm_indic
    from anuvani_diar_ref.decoder import load_decoder
    manifest = json.loads((a.bundle / "MANIFEST.json").read_text())
    mode = manifest["modes"]["indic"]
    model = load_bundle(a.bundle, a.device, a.batch_size, a.threads)
    decoder = load_decoder(a.bundle / mode["decoder"])
    a.out.mkdir(parents=True, exist_ok=True)
    log = dict(version=manifest["version"], model_config=mode, bundle_sha256=manifest["files"],
               device=a.device, batch_size=a.batch_size, threads=a.threads, python=platform.python_version(), files=[])
    for f in files:
        t0 = time.monotonic()
        audio, sr = load_audio(f)
        d = analyse(model, audio, sr)
        uri = f.stem
        out = a.out / f"{uri}.rttm"
        write_rttm_indic(out, uri, indic_mode(d, mode["clustering"], decoder))
        rec = dict(input=str(f), rttm=str(out), duration_s=round(len(audio) / sr, 3), seconds=round(time.monotonic() - t0, 2),
                   rttm_sha256=hashlib.sha256(out.read_bytes()).hexdigest())
        log["files"].append(rec)
        print(json.dumps(rec), flush=True)
    (a.out / "run.json").write_text(json.dumps(log, indent=1))


if __name__ == "__main__":
    main()