Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
38 lines
1.2 KiB
Python
38 lines
1.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Run the official Smart Turn v3 ONNX inference path on a 16 kHz mono WAV."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
|
|
import numpy as np
|
|
import onnxruntime as ort
|
|
import soundfile as sf
|
|
from transformers import WhisperFeatureExtractor
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--onnx", required=True)
|
|
parser.add_argument("--audio", required=True)
|
|
args = parser.parse_args()
|
|
|
|
audio, rate = sf.read(args.audio, dtype="float32")
|
|
if rate != 16000 or audio.ndim != 1:
|
|
raise ValueError("expected 16 kHz mono audio")
|
|
audio = audio[-128000:]
|
|
if len(audio) < 128000:
|
|
audio = np.pad(audio, (128000 - len(audio), 0))
|
|
features = WhisperFeatureExtractor(chunk_length=8)(
|
|
audio, sampling_rate=16000, return_tensors="np", padding="max_length",
|
|
max_length=128000, truncation=True, do_normalize=True,
|
|
).input_features.astype(np.float32)
|
|
probability = float(ort.InferenceSession(args.onnx).run(
|
|
None, {"input_features": features}
|
|
)[0][0].item())
|
|
print(f"turn_complete={int(probability > 0.5)}")
|
|
print(f"turn_probability={probability:.9f}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|