| |
| """Nemotron-3-Diarization streaming diarization on AX650N: 16 kHz wav -> RTTM. |
| |
| python3 example.py meeting.wav out.rttm |
| python3 example.py meeting.wav out.rttm --chunk-len 6 --chunk-right-context 4 \ |
| --fifo-len 264 --spkcache-update-period 222 # 1.04 s low latency |
| """ |
|
|
| import argparse |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import soundfile as sf |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent)) |
|
|
| from nemotron_sdk import AxengineGraphPair, SortformerConfig, StreamingDiarizer, load_learnable_sil_emb |
|
|
| MODELS = Path(__file__).resolve().parents[1] / "models" |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("wav") |
| parser.add_argument("rttm", nargs="?", default=None) |
| parser.add_argument("--preencode", default=str(MODELS / "preencode.axmodel")) |
| parser.add_argument("--encoder", default=str(MODELS / "encoder.axmodel")) |
| parser.add_argument("--sil-emb", default=str(MODELS / "learnable_sil_emb.npy")) |
| parser.add_argument("--assets", default=str(MODELS), help="dir with feature_window.npy / mel_filterbank.npy") |
| parser.add_argument("--chunk-len", type=int, default=340) |
| parser.add_argument("--chunk-left-context", type=int, default=0) |
| parser.add_argument("--chunk-right-context", type=int, default=40) |
| parser.add_argument("--fifo-len", type=int, default=40) |
| parser.add_argument("--spkcache-len", type=int, default=264) |
| parser.add_argument("--spkcache-update-period", type=int, default=300) |
| parser.add_argument("--threads", type=int, default=8, help="reserved (numpy front-end is single-process)") |
| args = parser.parse_args() |
|
|
| config = SortformerConfig( |
| chunk_len=args.chunk_len, |
| chunk_left_context=args.chunk_left_context, |
| chunk_right_context=args.chunk_right_context, |
| fifo_len=args.fifo_len, |
| spkcache_len=args.spkcache_len, |
| spkcache_update_period=args.spkcache_update_period, |
| ) |
| if not Path(args.sil_emb).exists(): |
| raise FileNotFoundError(f"learnable silence embedding not found: {args.sil_emb}") |
| load_learnable_sil_emb(config, args.sil_emb) |
|
|
| waveform, sample_rate = sf.read(args.wav, dtype="float32", always_2d=True) |
| waveform = waveform.mean(axis=1) |
| if sample_rate != 16000: |
| raise SystemExit(f"expected 16 kHz wav, got {sample_rate}") |
|
|
| diarizer = StreamingDiarizer( |
| AxengineGraphPair(args.preencode, args.encoder), config, feature_assets=args.assets |
| ) |
| start = time.perf_counter() |
| preds = diarizer.process_wav(waveform, sample_rate) |
| elapsed = time.perf_counter() - start |
| duration = waveform.shape[0] / sample_rate |
|
|
| uri = Path(args.wav).stem |
| lines = diarizer.rttm_lines(preds, uri) |
| if args.rttm: |
| Path(args.rttm).write_text("\n".join(lines) + "\n") |
|
|
| print( |
| f"{uri}: {len(preds)} frames, {len(lines)} segments, " |
| f"{elapsed:.1f}s / {duration:.1f}s audio (RTF {elapsed / duration:.4f})" |
| ) |
| if args.rttm: |
| print(f"rttm -> {args.rttm}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
