| |
| |
| """anuvani-diar: the release's `diarize.py --mode indic` command line, on the Rust engine. |
| |
| anuvani-diar --out OUT_DIR audio1.wav [audio2.flac ...] |
| anuvani-diar --out OUT_DIR --in-dir AUDIO_DIR |
| |
| RTTM files byte-identical to release/anuvani_diar_v1/diarize.py --mode indic, and its other flags. Differences: there |
| is one model, so no --mode; run.json names the model's settings `model_config`; every visible GPU is used unless |
| --devices says otherwise; files finish in any order (run.json keeps the input order); a file that fails gets an "error" |
| entry instead of stopping the run, and the exit status is 1. |
| """ |
| import argparse |
| import hashlib |
| import json |
| import platform |
| import sys |
| from pathlib import Path |
|
|
| AUDIO_EXT = ("*.wav", "*.flac", "*.mp3", "*.m4a", "*.ogg") |
|
|
|
|
| def main(argv=None): |
| 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=None) |
| p.add_argument("--devices", "--device", dest="devices", default="auto", help="which devices run the networks: auto (every visible GPU, else the CPU), a number of GPUs (2: the first two), GPU indices (0,3), 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(argv) |
| if a.batch_size != 8: |
| raise SystemExit("--batch-size: only 8, the release's verified setting, is supported") |
| 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 .api import Diarizer, devices_arg |
| d = Diarizer(bundle=a.bundle, devices=devices_arg(a.devices), threads=a.threads) |
| manifest = d.manifest |
| a.out.mkdir(parents=True, exist_ok=True) |
| log = dict(version=manifest["version"], model_config=manifest["modes"]["indic"], bundle_sha256=manifest["files"], |
| devices=a.devices, batch_size=a.batch_size, threads=a.threads, python=platform.python_version(), files=[]) |
| recs, failed = [None] * len(files), 0 |
| for r in d.run(files): |
| f = files[r.index] |
| if r.ok: |
| out = r.save_rttm(a.out / f"{r.uri}.rttm") |
| rec = dict(input=str(f), rttm=str(out), duration_s=round(r.duration_s, 3), seconds=round(r.timing["wall_s"], 2), |
| rttm_sha256=hashlib.sha256(out.read_bytes()).hexdigest()) |
| else: |
| failed += 1 |
| rec = dict(input=str(f), error=r.error) |
| recs[r.index] = rec |
| print(json.dumps(rec), flush=True) |
| log["files"] = recs |
| (a.out / "run.json").write_text(json.dumps(log, indent=1)) |
| return 1 if failed else 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
