Fully functional StreamVAD

.gitignore

@@ -0,0 +1,2 @@
+output
+__pycache__

README.md

@@ -1,3 +1,46 @@
 ---
 license: mit
 ---
+# SileroVAD
+
+流式语音端点识别
+
+
+## Demo
+
+```
+python main.py --input demo.wav --output_dir output --model silero_vad.onnx
+```
+
+被分段的语音后保存在output目录中
+
+
+## 在项目中使用
+
+1. 复制silero_vad.onnx SileroOrt.py StreamVAD.py 三个文件到项目中
+2. from StreamVAD import StreamVAD
+3.
+初始化 
+```
+vad = StreamVAD(args.model, 
+                    sensitivity=0.5,
+                    silence_ms=200)
+```
+
+运行
+```
+for result in vad.run(audio, vad.model.sr):
+    if result:
+        print(result)
+```
+
+result的格式为:
+```
+{
+    'start_ts': 语音开始的时间
+    'end_ts': 语音结束的时间
+    'audio': 语音数据
+}
+```
+
+时间戳的格式可通过StreamVAD.datetime_format设置

SileroOrt.py

@@ -0,0 +1,71 @@
+import numpy as np
+import onnxruntime as ort
+import librosa
+
+
+class SileroOrt:
+    def __init__(self, model_path: str):
+        super().__init__()
+
+        self.batch_size = 1
+        self.sr = 16000
+        self.hidden_size = 128
+        self.context_size = 64 if self.sr == 16000 else 32
+        self.context = np.zeros((self.batch_size, self.context_size), dtype=np.float32)
+        self.state = np.zeros((2, self.batch_size, self.hidden_size), dtype=np.float32)
+        self.num_samples = 512 if self.sr == 16000 else 256
+
+        self.model = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
+        self.reset_states()
+
+
+    def reset_states(self):
+        self.context = np.zeros((self.batch_size, self.context_size), dtype=np.float32)
+        self.state = np.zeros((2, self.batch_size, self.hidden_size), dtype=np.float32)
+
+
+    def __call__(self, x):
+        if len(x.shape) == 1:
+            x = x[None, ...]
+
+        data = np.concatenate([self.context, x], axis=1)
+        data = np.pad(data, ((0, 0), (0, 64)), 'reflect')
+        input_feed = {
+            "data": data,
+            "state": self.state
+        }
+
+        output, self.state = self.model.run(None, input_feed=input_feed)
+        self.context = x[..., -self.context_size:]
+
+        if len(output.shape) == 0:
+            output = np.array([output], dtype=np.float32)
+        return output
+    
+
+    def audio_forward(self, x: np.ndarray, sr: int):
+        if len(x.shape) > 1:
+            # mono
+            x = x[0]
+
+        if x.dtype == np.int16:
+            x = x.astype(np.float32) / 32768
+
+        if sr != self.sr:
+            x = librosa.resample(x, orig_sr=sr, target_sr=self.sr)
+
+        outs = []
+        num_samples = self.num_samples
+
+        if x.shape[0] % num_samples:
+            pad_num = num_samples - (x.shape[0] % num_samples)
+            x = np.pad(x, ((0, pad_num)), 'constant', value=0.0)
+
+        for i in range(0, x.shape[0], num_samples):
+            wavs_batch = x[i:i+num_samples]
+            out_chunk = self.__call__(wavs_batch)
+            # print(out_chunk)
+            outs.append(out_chunk)
+
+        stacked = np.concatenate(outs, axis=-1)
+        return stacked

StreamVAD.py

@@ -0,0 +1,78 @@
+from SileroOrt import SileroOrt
+import numpy as np
+from datetime import datetime, timedelta
+
+
+class StreamVAD:
+    def __init__(self, model_path, 
+                        sensitivity=0.5,
+                        silence_ms=200,
+                        datetime_format='%Y-%m-%d %H:%M:%S.%f'):
+        '''
+        model_path: path of silero_vad.onnx
+        sensitivity: thresh of voice activation, 
+            higher means more sensitive, 
+            hence, low speech prob thresh
+        silence_ms: pop audio after silence for silence_ms milliseconds
+        datetime_format: format of datetime in return data
+        '''
+
+        self.model = SileroOrt(model_path)
+        self.sensitivity = sensitivity
+        self.silence_ms = silence_ms
+        self.datetime_format = datetime_format
+
+        self.reset()
+
+
+    def reset(self):
+        self.silence_count = 0
+        self.speech_count = 0
+        self.return_data = {
+            "start_ts": '',
+            "end_ts": '',
+            "audio": None
+        }
+        self.vad_data_list = []
+        self.model.reset_states()
+
+    
+    def run(self, audio: np.ndarray, sr: int = 16000):
+        # record datetime
+        cur_ts = datetime.now()
+
+        # inference
+        speech_probs = self.model.audio_forward(audio, sr)[0]
+
+        for i, prob in enumerate(speech_probs):
+            audio_slice = audio[i * self.model.num_samples : (i + 1) * self.model.num_samples]
+            ts = cur_ts.strftime(self.datetime_format)
+
+            # is speech
+            if prob > 1 - self.sensitivity:
+                self.silence_count = 0
+                # new speech segment
+                if self.speech_count == 0:
+                    self.return_data['start_ts'] = ts
+
+                self.speech_count += 1
+                self.vad_data_list.append(audio_slice)
+            # silence
+            else:
+                if self.speech_count > 0:
+                    self.silence_count += 1
+
+                    # exceed silence limit
+                    if 1000 * self.silence_count * self.model.num_samples / self.model.sr > self.silence_ms:
+                        # return audio segment
+                        self.return_data['end_ts'] = ts
+                        self.return_data['audio'] = np.concatenate(self.vad_data_list, axis=-1)
+
+                        yield self.return_data
+
+                        self.reset()
+                    else:
+                        self.vad_data_list.append(audio_slice)
+
+            # timestamp
+            cur_ts += timedelta(seconds=self.model.num_samples / self.model.sr)

main.py

@@ -0,0 +1,35 @@
+import argparse
+from StreamVAD import StreamVAD
+import os
+import librosa
+import soundfile as sf
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--input', type=str, required=True, help='Input audio file')
+    parser.add_argument('--model', type=str, default='./silero_vad.onnx')
+    parser.add_argument('--output_dir', type=str, default='output', help='Output audio dir')
+    return parser.parse_args()
+
+
+def main():
+    args = get_args()
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    vad = StreamVAD(args.model, 
+                    sensitivity=0.5,
+                    silence_ms=200)
+    
+    audio, _ = librosa.load(args.input, sr=vad.model.sr, mono=True)
+    i = 0
+    for result in vad.run(audio, vad.model.sr):
+        if result:
+            print(result)
+            filename = os.path.join(args.output_dir, f"{i}.wav")
+            sf.write(filename, result['audio'], samplerate=vad.model.sr)
+            i += 1
+
+
+if __name__ == '__main__':
+    main()

requirements.txt

@@ -0,0 +1,6 @@
+onnxruntime==1.17.0
+librosa
+numpy<2.0
+samplerate
+resampy
+soundfile

silero_vad.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:144f7a8e8db2bbe7e90407f966ec811cbcdc7258fffbc867798597a33c957118
+size 1247953