custom_stream.py¶

"""Custom audio stream for the TrulyNatural SDK Python binding.

Demonstrates ``Stream.from_provider``: build a ``snsr.Stream``
from Python callbacks and hand it to a session via
``Session.set_stream``. This is the right shape when the audio
does not live in a regular file -- for example a network socket,
an in-memory ring buffer, or a hardware abstraction that does
not look like a microphone.

Here we keep the example small by reading the same WAV file as
``hello_world.py`` into memory once, stripping the 44-byte RIFF
WAVE header, and then serving the raw PCM through a Python
``read`` callback. The SDK pulls one chunk at a time on its
processing thread; the callback fills the buffer and returns the
number of bytes written, returning a value smaller than the
buffer to signal end-of-stream.

Modelled on ``tests/test_snsr.py::test_Stream_from_provider_read``.

Usage::

    uv run src/custom_stream.py [--sdk-root PATH]
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

import snsr


MODEL = "spot-voicegenie-enUS-6.5.1-m.snsr"
AUDIO = "voice-genie-set-cruise-control.wav"
RIFF_HEADER_BYTES = 44


def default_sdk_root() -> Path:
    return Path(__file__).resolve().parents[3]


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument(
        "--sdk-root",
        type=Path,
        default=default_sdk_root(),
        help="TrulyNatural SDK install root (default: auto-detect)",
    )
    return parser.parse_args(argv)


def make_pcm_provider(pcm: bytes) -> snsr.Stream:
    """Return a snsr.Stream that yields ``pcm`` over ``read`` callbacks."""
    offset = 0

    def read(buffer: bytearray) -> int:
        nonlocal offset
        n = min(len(buffer), len(pcm) - offset)
        buffer[:n] = pcm[offset : offset + n]
        offset += n
        return n

    return snsr.Stream.from_provider(read=read)


def run_custom_stream(model_path: Path, audio_path: Path) -> int:
    """Run the spotter on ``audio_path`` via a from_provider Stream."""
    pcm = audio_path.read_bytes()[RIFF_HEADER_BYTES:]
    count = 0

    def on_result(s: snsr.Session, _key: bytes) -> None:
        nonlocal count
        count += 1
        text = s.get_string(snsr.RES_TEXT)
        begin_ms = s.get_int(snsr.RES_BEGIN_MS)
        end_ms = s.get_double(snsr.RES_END_MS)
        score = s.get_double(snsr.RES_SCORE)
        print(f"  spotted {text!r}: {begin_ms:>5} ms - {end_ms:>7.1f} ms (score {score:.4f})")

    print(f"snsr {snsr.VERSION}")
    print(f"  model: {model_path}")
    print(f"  audio: {audio_path} ({len(pcm)} bytes of PCM via from_provider)")
    print()

    with snsr.Session(str(model_path)) as s:
        s.require(snsr.TASK_TYPE, snsr.PHRASESPOT)
        s.set_handler(snsr.RESULT_EVENT, on_result)
        with make_pcm_provider(pcm) as audio:
            s.set_stream(snsr.SOURCE_AUDIO_PCM, audio)
            s.run()

    print()
    print(f"done: {count} result event(s)")
    return count


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    sdk_root: Path = args.sdk_root.resolve()

    model_path = sdk_root / "model" / MODEL
    audio_path = sdk_root / "data" / "audio" / AUDIO

    for label, path in (("model", model_path), ("audio", audio_path)):
        if not path.is_file():
            print(f"error: {label} not found: {path}", file=sys.stderr)
            print(
                f"hint: pass --sdk-root pointing at a TrulyNatural SDK install",
                file=sys.stderr,
            )
            return 2

    run_custom_stream(model_path, audio_path)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())