222 lines
7.9 KiB
Python
222 lines
7.9 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
from abc import ABC, abstractmethod
|
|
from dataclasses import dataclass
|
|
from typing import Any, Callable
|
|
|
|
|
|
LOGGER = logging.getLogger("uvicorn.error")
|
|
|
|
|
|
@dataclass
|
|
class VADFrameResult:
|
|
is_speech: bool = False
|
|
speech_started: bool = False
|
|
speech_ended: bool = False
|
|
speech_probability: float = 0.0
|
|
utterance_audio: bytes | None = None
|
|
|
|
|
|
class BaseVAD(ABC):
|
|
@abstractmethod
|
|
def feed(self, audio_chunk: bytes) -> VADFrameResult:
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def flush(self) -> bytes | None:
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def reset(self) -> None:
|
|
raise NotImplementedError
|
|
|
|
|
|
class SileroVADDetector(BaseVAD):
|
|
def __init__(
|
|
self,
|
|
*,
|
|
sample_rate_hz: int = 8000,
|
|
threshold: float = 0.5,
|
|
negative_threshold: float | None = None,
|
|
speech_end_silence_ms: int = 1600,
|
|
speech_pad_ms: int = 64,
|
|
min_speech_duration_ms: int = 0,
|
|
model: Any | None = None,
|
|
prediction_fn: Callable[[bytes], float] | None = None,
|
|
use_onnx: bool = False,
|
|
) -> None:
|
|
if sample_rate_hz not in {8000, 16000}:
|
|
raise ValueError("Silero VAD supports only 8000 Hz and 16000 Hz sample rates")
|
|
self.sample_rate_hz = sample_rate_hz
|
|
self.threshold = max(min(threshold, 1.0), 0.0)
|
|
self.negative_threshold = (
|
|
max(min(negative_threshold, 1.0), 0.0)
|
|
if negative_threshold is not None
|
|
else max(self.threshold - 0.15, 0.01)
|
|
)
|
|
self.speech_end_silence_ms = max(speech_end_silence_ms, 32)
|
|
self.speech_pad_ms = max(speech_pad_ms, 0)
|
|
self.min_speech_duration_ms = max(min_speech_duration_ms, 0)
|
|
self.window_samples = 512 if self.sample_rate_hz == 16000 else 256
|
|
self.window_bytes = self.window_samples * 2
|
|
self._speech_end_silence_samples = int(self.sample_rate_hz * self.speech_end_silence_ms / 1000.0)
|
|
self._speech_pad_bytes = int(self.sample_rate_hz * self.speech_pad_ms / 1000.0) * 2
|
|
self._min_speech_samples = int(self.sample_rate_hz * self.min_speech_duration_ms / 1000.0)
|
|
self._model = model
|
|
self._prediction_fn = prediction_fn
|
|
self._use_onnx = use_onnx
|
|
self.reset()
|
|
|
|
def reset(self) -> None:
|
|
self._window_buffer = bytearray()
|
|
self._pre_speech_audio = bytearray()
|
|
self._utterance_audio = bytearray()
|
|
self._triggered = False
|
|
self._silence_samples = 0
|
|
self._speech_samples = 0
|
|
if self._model is not None and hasattr(self._model, "reset_states"):
|
|
self._model.reset_states()
|
|
|
|
def feed(self, audio_chunk: bytes) -> VADFrameResult:
|
|
result = VADFrameResult()
|
|
if not audio_chunk:
|
|
return result
|
|
|
|
self._window_buffer.extend(audio_chunk)
|
|
while len(self._window_buffer) >= self.window_bytes:
|
|
window = bytes(self._window_buffer[: self.window_bytes])
|
|
del self._window_buffer[: self.window_bytes]
|
|
speech_probability = self._predict_speech_probability(window)
|
|
result.speech_probability = speech_probability
|
|
result.is_speech = result.is_speech or self._triggered or speech_probability >= self.threshold
|
|
|
|
if not self._triggered:
|
|
if speech_probability >= self.threshold:
|
|
self._triggered = True
|
|
self._speech_samples = self.window_samples
|
|
self._silence_samples = 0
|
|
self._utterance_audio = bytearray(self._pre_speech_audio)
|
|
self._utterance_audio.extend(window)
|
|
self._pre_speech_audio.clear()
|
|
result.speech_started = True
|
|
result.is_speech = True
|
|
else:
|
|
self._append_pre_speech_window(window)
|
|
continue
|
|
|
|
self._utterance_audio.extend(window)
|
|
result.is_speech = True
|
|
|
|
if speech_probability >= self.threshold:
|
|
self._silence_samples = 0
|
|
self._speech_samples += self.window_samples
|
|
continue
|
|
|
|
if speech_probability >= self.negative_threshold:
|
|
self._silence_samples = 0
|
|
continue
|
|
|
|
self._silence_samples += self.window_samples
|
|
if self._silence_samples < self._speech_end_silence_samples:
|
|
continue
|
|
|
|
utterance_audio = bytes(self._utterance_audio)
|
|
speech_samples = self._speech_samples
|
|
self._reset_segment()
|
|
result.is_speech = False
|
|
if speech_samples >= self._min_speech_samples:
|
|
result.speech_ended = True
|
|
result.utterance_audio = utterance_audio
|
|
|
|
return result
|
|
|
|
def flush(self) -> bytes | None:
|
|
if not self._triggered or not self._utterance_audio:
|
|
return None
|
|
utterance_audio = bytes(self._utterance_audio)
|
|
speech_samples = self._speech_samples
|
|
self._reset_segment()
|
|
if speech_samples < self._min_speech_samples:
|
|
return None
|
|
return utterance_audio
|
|
|
|
def _reset_segment(self) -> None:
|
|
self._triggered = False
|
|
self._silence_samples = 0
|
|
self._speech_samples = 0
|
|
self._utterance_audio = bytearray()
|
|
self._pre_speech_audio.clear()
|
|
|
|
def _append_pre_speech_window(self, window: bytes) -> None:
|
|
self._pre_speech_audio.extend(window)
|
|
if self._speech_pad_bytes <= 0:
|
|
self._pre_speech_audio.clear()
|
|
return
|
|
overflow = len(self._pre_speech_audio) - self._speech_pad_bytes
|
|
if overflow > 0:
|
|
del self._pre_speech_audio[:overflow]
|
|
|
|
def _predict_speech_probability(self, window: bytes) -> float:
|
|
if self._prediction_fn is not None:
|
|
return max(0.0, min(float(self._prediction_fn(window)), 1.0))
|
|
|
|
model = self._load_model()
|
|
try:
|
|
import numpy as np
|
|
import torch
|
|
except Exception as exc: # noqa: BLE001
|
|
raise RuntimeError(
|
|
"Silero VAD requires numpy and torch at runtime; add them to requirements.txt"
|
|
) from exc
|
|
|
|
pcm = np.frombuffer(window, dtype=np.int16).astype(np.float32) / 32768.0
|
|
tensor = torch.from_numpy(pcm)
|
|
probability = model(tensor, self.sample_rate_hz).item()
|
|
return max(0.0, min(float(probability), 1.0))
|
|
|
|
def _load_model(self) -> Any:
|
|
if self._model is not None:
|
|
return self._model
|
|
try:
|
|
import torch
|
|
except Exception as exc: # noqa: BLE001
|
|
raise RuntimeError(
|
|
"Silero VAD requires torch; install torch/torchaudio or use the ONNX runtime path"
|
|
) from exc
|
|
|
|
torch.set_num_threads(1)
|
|
try:
|
|
from silero_vad import load_silero_vad
|
|
except Exception: # noqa: BLE001
|
|
load_silero_vad = None
|
|
|
|
if load_silero_vad is not None:
|
|
self._model = load_silero_vad(onnx=self._use_onnx)
|
|
LOGGER.info(
|
|
"realtime voice VAD loaded via silero_vad package sample_rate=%s onnx=%s",
|
|
self.sample_rate_hz,
|
|
self._use_onnx,
|
|
)
|
|
return self._model
|
|
|
|
try:
|
|
self._model, _ = torch.hub.load(
|
|
repo_or_dir="snakers4/silero-vad",
|
|
model="silero_vad",
|
|
onnx=self._use_onnx,
|
|
force_reload=False,
|
|
)
|
|
except TypeError:
|
|
self._model, _ = torch.hub.load(
|
|
repo_or_dir="snakers4/silero-vad",
|
|
model="silero_vad",
|
|
force_reload=False,
|
|
)
|
|
LOGGER.info(
|
|
"realtime voice VAD loaded via torch.hub sample_rate=%s onnx=%s",
|
|
self.sample_rate_hz,
|
|
self._use_onnx,
|
|
)
|
|
return self._model
|