File size: 2,149 Bytes
7c268e9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | #!/usr/bin/env python3
"""Sortformer streaming diarization on AX650N: 16 kHz wav -> RTTM.
python3 example.py meeting.wav out.rttm [--fast]
"""
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 sortformer_sdk import AxengineGraphPair, SortformerConfig, StreamingDiarizer # noqa: E402
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("--fast", action="store_true", help="fifo40 encoder: RTF 0.064, DER +0.3~0.6pp")
parser.add_argument("--preencode", default=str(MODELS / "preencode.axmodel"))
parser.add_argument("--encoder", default=None)
parser.add_argument("--threads", type=int, default=8, help="reserved (numpy front-end is single-process)")
args = parser.parse_args()
encoder = args.encoder or str(MODELS / ("encoder_fifo40.axmodel" if args.fast else "encoder.axmodel"))
config = SortformerConfig(
fifo_len=40 if args.fast else 188,
spkcache_update_period=31 if args.fast else 144,
)
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, encoder), config)
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()
|