| |
| |
| |
| """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() |
|
|
