|
import torch |
|
|
|
|
|
class VADIterator: |
|
def __init__( |
|
self, |
|
model, |
|
threshold: float = 0.5, |
|
sampling_rate: int = 16000, |
|
min_silence_duration_ms: int = 100, |
|
speech_pad_ms: int = 30, |
|
): |
|
""" |
|
Mainly taken from https://github.com/snakers4/silero-vad |
|
Class for stream imitation |
|
|
|
Parameters |
|
---------- |
|
model: preloaded .jit/.onnx silero VAD model |
|
|
|
threshold: float (default - 0.5) |
|
Speech threshold. Silero VAD outputs speech probabilities for each audio chunk, probabilities ABOVE this value are considered as SPEECH. |
|
It is better to tune this parameter for each dataset separately, but "lazy" 0.5 is pretty good for most datasets. |
|
|
|
sampling_rate: int (default - 16000) |
|
Currently silero VAD models support 8000 and 16000 sample rates |
|
|
|
min_silence_duration_ms: int (default - 100 milliseconds) |
|
In the end of each speech chunk wait for min_silence_duration_ms before separating it |
|
|
|
speech_pad_ms: int (default - 30 milliseconds) |
|
Final speech chunks are padded by speech_pad_ms each side |
|
""" |
|
|
|
self.model = model |
|
self.threshold = threshold |
|
self.sampling_rate = sampling_rate |
|
self.is_speaking = False |
|
self.buffer = [] |
|
|
|
if sampling_rate not in [8000, 16000]: |
|
raise ValueError( |
|
"VADIterator does not support sampling rates other than [8000, 16000]" |
|
) |
|
|
|
self.min_silence_samples = sampling_rate * min_silence_duration_ms / 1000 |
|
self.speech_pad_samples = sampling_rate * speech_pad_ms / 1000 |
|
self.reset_states() |
|
|
|
def reset_states(self): |
|
self.model.reset_states() |
|
self.triggered = False |
|
self.temp_end = 0 |
|
self.current_sample = 0 |
|
|
|
@torch.no_grad() |
|
def __call__(self, x): |
|
""" |
|
x: torch.Tensor |
|
audio chunk (see examples in repo) |
|
|
|
return_seconds: bool (default - False) |
|
whether return timestamps in seconds (default - samples) |
|
""" |
|
|
|
if not torch.is_tensor(x): |
|
try: |
|
x = torch.Tensor(x) |
|
except Exception: |
|
raise TypeError("Audio cannot be casted to tensor. Cast it manually") |
|
|
|
window_size_samples = len(x[0]) if x.dim() == 2 else len(x) |
|
self.current_sample += window_size_samples |
|
|
|
speech_prob = self.model(x, self.sampling_rate).item() |
|
|
|
if (speech_prob >= self.threshold) and self.temp_end: |
|
self.temp_end = 0 |
|
|
|
if (speech_prob >= self.threshold) and not self.triggered: |
|
self.triggered = True |
|
return None |
|
|
|
if (speech_prob < self.threshold - 0.15) and self.triggered: |
|
if not self.temp_end: |
|
self.temp_end = self.current_sample |
|
if self.current_sample - self.temp_end < self.min_silence_samples: |
|
return None |
|
else: |
|
|
|
self.temp_end = 0 |
|
self.triggered = False |
|
spoken_utterance = self.buffer |
|
self.buffer = [] |
|
return spoken_utterance |
|
|
|
if self.triggered: |
|
self.buffer.append(x) |
|
|
|
return None |
|
|
