This commit is contained in:
@@ -0,0 +1,295 @@
|
||||
"""
|
||||
Voice Activity Detection (VAD) processor using Silero-VAD (ONNX version).
|
||||
|
||||
Detects speech segments in real-time audio streams.
|
||||
Uses ONNX runtime to avoid torchaudio dependency issues.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
from typing import List, Optional, Tuple
|
||||
from dataclasses import dataclass
|
||||
import threading
|
||||
import os
|
||||
import urllib.request
|
||||
|
||||
from .models import SpeechSegment, VADEvent, VADEventType, MessageType
|
||||
|
||||
|
||||
@dataclass
|
||||
class VADState:
|
||||
"""Internal state of the VAD processor."""
|
||||
is_speech_active: bool = False
|
||||
speech_start_sample: int = 0
|
||||
silence_start_sample: int = 0
|
||||
last_speech_prob: float = 0.0
|
||||
total_samples_processed: int = 0
|
||||
|
||||
|
||||
class VADProcessor:
|
||||
"""
|
||||
Voice Activity Detection using Silero-VAD (ONNX version).
|
||||
|
||||
Features:
|
||||
- Real-time speech detection
|
||||
- Configurable thresholds for speech/silence duration
|
||||
- Event generation for speech start/end
|
||||
- Thread-safe operations
|
||||
- No torchaudio dependency (uses ONNX runtime)
|
||||
"""
|
||||
|
||||
# Silero VAD ONNX model URL
|
||||
ONNX_MODEL_URL = "https://github.com/snakers4/silero-vad/raw/master/src/silero_vad/data/silero_vad.onnx"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
sample_rate: int = 16000,
|
||||
threshold: float = 0.5,
|
||||
min_speech_duration_ms: int = 250,
|
||||
min_silence_duration_ms: int = 500,
|
||||
window_size_samples: int = 512,
|
||||
min_volume_threshold: float = 0.01,
|
||||
):
|
||||
"""
|
||||
Initialize the VAD processor.
|
||||
|
||||
Args:
|
||||
sample_rate: Audio sample rate (must be 16000 for Silero-VAD)
|
||||
threshold: Speech probability threshold (0.0-1.0)
|
||||
min_speech_duration_ms: Minimum speech duration to trigger speech_start
|
||||
min_silence_duration_ms: Minimum silence duration to trigger speech_end
|
||||
window_size_samples: VAD window size (512 for 16kHz = 32ms)
|
||||
min_volume_threshold: Minimum RMS volume (0.0-1.0) to consider as potential speech
|
||||
"""
|
||||
if sample_rate != 16000:
|
||||
raise ValueError("Silero-VAD requires 16kHz sample rate")
|
||||
|
||||
self.sample_rate = sample_rate
|
||||
self.threshold = threshold
|
||||
self.min_speech_samples = int(min_speech_duration_ms * sample_rate / 1000)
|
||||
self.min_silence_samples = int(min_silence_duration_ms * sample_rate / 1000)
|
||||
self.window_size = window_size_samples
|
||||
self.min_volume_threshold = min_volume_threshold
|
||||
|
||||
# Load ONNX model
|
||||
self._session = None
|
||||
self._load_model()
|
||||
|
||||
# ONNX model state - single state tensor (size depends on model version)
|
||||
# Silero VAD v5 uses a single 'state' tensor of shape (2, 1, 128)
|
||||
self._state_tensor = np.zeros((2, 1, 128), dtype=np.float32)
|
||||
|
||||
# State
|
||||
self._state = VADState()
|
||||
self._lock = threading.Lock()
|
||||
|
||||
# Pending speech segment (being accumulated)
|
||||
self._pending_segment_start: Optional[int] = None
|
||||
|
||||
def _get_model_path(self) -> str:
|
||||
"""Get path to ONNX model, downloading if necessary."""
|
||||
cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "silero-vad")
|
||||
os.makedirs(cache_dir, exist_ok=True)
|
||||
model_path = os.path.join(cache_dir, "silero_vad.onnx")
|
||||
|
||||
if not os.path.exists(model_path):
|
||||
print(f"Downloading Silero-VAD ONNX model to {model_path}...")
|
||||
urllib.request.urlretrieve(self.ONNX_MODEL_URL, model_path)
|
||||
print("Download complete.")
|
||||
|
||||
return model_path
|
||||
|
||||
def _load_model(self) -> None:
|
||||
"""Load Silero-VAD ONNX model."""
|
||||
try:
|
||||
import onnxruntime as ort
|
||||
|
||||
model_path = self._get_model_path()
|
||||
self._session = ort.InferenceSession(
|
||||
model_path,
|
||||
providers=['CPUExecutionProvider']
|
||||
)
|
||||
print(f"Silero-VAD ONNX model loaded from {model_path}")
|
||||
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to load Silero-VAD ONNX model: {e}")
|
||||
|
||||
def _run_inference(self, audio_window: np.ndarray) -> float:
|
||||
"""Run VAD inference on a single window."""
|
||||
# Prepare input
|
||||
audio_input = audio_window.reshape(1, -1).astype(np.float32)
|
||||
sr_input = np.array([self.sample_rate], dtype=np.int64)
|
||||
|
||||
# Run inference - Silero VAD v5 uses 'state' instead of 'h'/'c'
|
||||
outputs = self._session.run(
|
||||
['output', 'stateN'],
|
||||
{
|
||||
'input': audio_input,
|
||||
'sr': sr_input,
|
||||
'state': self._state_tensor,
|
||||
}
|
||||
)
|
||||
|
||||
# Update state
|
||||
speech_prob = outputs[0][0][0]
|
||||
self._state_tensor = outputs[1]
|
||||
|
||||
return float(speech_prob)
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset VAD state for a new session."""
|
||||
with self._lock:
|
||||
self._state = VADState()
|
||||
self._pending_segment_start = None
|
||||
# Reset state tensor
|
||||
self._state_tensor = np.zeros((2, 1, 128), dtype=np.float32)
|
||||
|
||||
def process(
|
||||
self,
|
||||
audio_chunk: np.ndarray,
|
||||
return_events: bool = True,
|
||||
) -> Tuple[List[SpeechSegment], List[VADEvent]]:
|
||||
"""
|
||||
Process an audio chunk and detect speech segments.
|
||||
|
||||
Args:
|
||||
audio_chunk: Audio data as float32 array
|
||||
return_events: Whether to return VAD events
|
||||
|
||||
Returns:
|
||||
Tuple of (completed_segments, events)
|
||||
"""
|
||||
if audio_chunk.dtype != np.float32:
|
||||
audio_chunk = audio_chunk.astype(np.float32)
|
||||
|
||||
completed_segments: List[SpeechSegment] = []
|
||||
events: List[VADEvent] = []
|
||||
|
||||
with self._lock:
|
||||
# Process in windows
|
||||
chunk_start_sample = self._state.total_samples_processed
|
||||
num_windows = len(audio_chunk) // self.window_size
|
||||
|
||||
for i in range(num_windows):
|
||||
window_start = i * self.window_size
|
||||
window_end = window_start + self.window_size
|
||||
window = audio_chunk[window_start:window_end]
|
||||
|
||||
# Check volume (RMS) threshold first
|
||||
rms = np.sqrt(np.mean(window ** 2))
|
||||
if rms < self.min_volume_threshold:
|
||||
# Volume too low, treat as silence
|
||||
speech_prob = 0.0
|
||||
else:
|
||||
# Get speech probability from VAD model
|
||||
speech_prob = self._run_inference(window)
|
||||
|
||||
self._state.last_speech_prob = speech_prob
|
||||
|
||||
current_sample = chunk_start_sample + window_end
|
||||
is_speech = speech_prob >= self.threshold
|
||||
|
||||
# State machine for speech detection
|
||||
if is_speech:
|
||||
if not self._state.is_speech_active:
|
||||
# Potential speech start
|
||||
if self._pending_segment_start is None:
|
||||
self._pending_segment_start = current_sample - self.window_size
|
||||
|
||||
# Check if speech duration exceeds minimum
|
||||
speech_duration = current_sample - self._pending_segment_start
|
||||
if speech_duration >= self.min_speech_samples:
|
||||
self._state.is_speech_active = True
|
||||
self._state.speech_start_sample = self._pending_segment_start
|
||||
|
||||
if return_events:
|
||||
events.append(VADEvent(
|
||||
type=MessageType.VAD_EVENT,
|
||||
event=VADEventType.SPEECH_START,
|
||||
audio_timestamp_sec=self._pending_segment_start / self.sample_rate,
|
||||
))
|
||||
else:
|
||||
# Continue speech, reset silence counter
|
||||
self._state.silence_start_sample = 0
|
||||
else:
|
||||
if self._state.is_speech_active:
|
||||
# Potential speech end
|
||||
if self._state.silence_start_sample == 0:
|
||||
self._state.silence_start_sample = current_sample
|
||||
|
||||
# Check if silence duration exceeds minimum
|
||||
silence_duration = current_sample - self._state.silence_start_sample
|
||||
if silence_duration >= self.min_silence_samples:
|
||||
# Speech ended - create completed segment
|
||||
segment = SpeechSegment(
|
||||
start_sample=self._state.speech_start_sample,
|
||||
end_sample=self._state.silence_start_sample,
|
||||
start_sec=self._state.speech_start_sample / self.sample_rate,
|
||||
end_sec=self._state.silence_start_sample / self.sample_rate,
|
||||
)
|
||||
completed_segments.append(segment)
|
||||
|
||||
if return_events:
|
||||
events.append(VADEvent(
|
||||
type=MessageType.VAD_EVENT,
|
||||
event=VADEventType.SPEECH_END,
|
||||
audio_timestamp_sec=self._state.silence_start_sample / self.sample_rate,
|
||||
))
|
||||
|
||||
# Reset state
|
||||
self._state.is_speech_active = False
|
||||
self._state.speech_start_sample = 0
|
||||
self._state.silence_start_sample = 0
|
||||
self._pending_segment_start = None
|
||||
else:
|
||||
# No speech, reset pending
|
||||
self._pending_segment_start = None
|
||||
|
||||
# Update total samples processed
|
||||
self._state.total_samples_processed += len(audio_chunk)
|
||||
|
||||
return completed_segments, events
|
||||
|
||||
def force_end_speech(self) -> Optional[SpeechSegment]:
|
||||
"""
|
||||
Force end of current speech segment (e.g., when session ends).
|
||||
|
||||
Returns:
|
||||
Completed speech segment if speech was active, None otherwise
|
||||
"""
|
||||
with self._lock:
|
||||
if self._state.is_speech_active:
|
||||
segment = SpeechSegment(
|
||||
start_sample=self._state.speech_start_sample,
|
||||
end_sample=self._state.total_samples_processed,
|
||||
start_sec=self._state.speech_start_sample / self.sample_rate,
|
||||
end_sec=self._state.total_samples_processed / self.sample_rate,
|
||||
)
|
||||
|
||||
self._state.is_speech_active = False
|
||||
self._state.speech_start_sample = 0
|
||||
self._state.silence_start_sample = 0
|
||||
self._pending_segment_start = None
|
||||
|
||||
return segment
|
||||
|
||||
return None
|
||||
|
||||
@property
|
||||
def is_speech_active(self) -> bool:
|
||||
"""Check if speech is currently active."""
|
||||
with self._lock:
|
||||
return self._state.is_speech_active
|
||||
|
||||
@property
|
||||
def last_speech_probability(self) -> float:
|
||||
"""Get the last computed speech probability."""
|
||||
with self._lock:
|
||||
return self._state.last_speech_prob
|
||||
|
||||
@property
|
||||
def current_speech_duration_sec(self) -> float:
|
||||
"""Get duration of current speech segment (if active)."""
|
||||
with self._lock:
|
||||
if not self._state.is_speech_active:
|
||||
return 0.0
|
||||
return (self._state.total_samples_processed - self._state.speech_start_sample) / self.sample_rate
|
||||
Reference in New Issue
Block a user