python/example.py

#!/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()