STT_Server/tests/functor/vad_test.py

58 lines
1.6 KiB
Python

"""
Functor测试
VAD测试
"""
from src.functor.vad_functor import VADFunctor
from queue import Queue, Empty
from src.model_loader import ModelLoader
from src.models import AudioBinary_Config
from src.utils.data_format import wav_to_bytes
import time
from src.utils.logger import get_module_logger
logger = get_module_logger(__name__)
model_loader = ModelLoader()
def test_vad_functor():
# 加载模型
args = {
"vad_model": "fsmn-vad",
"vad_model_revision": "v2.0.4",
"auto_update": False,
}
model_loader.load_models(args)
# 创建VAD函数器
vad_functor = VADFunctor()
# 创建输入队列
input_queue = Queue()
# 设置输入队列
vad_functor.set_input_queue(input_queue)
# 设置音频配置
vad_functor.set_audio_config(AudioBinary_Config(
chunk_size=960,
chunk_stride=480,
sample_rate=16000,
sample_width=2,
channels=1
))
# 设置回调函数
vad_functor.add_callback(lambda x: print(x))
# 设置模型
vad_functor.set_model({
'vad': model_loader.models['vad']
})
# 启动VAD函数器
vad_functor.run()
# 加载数据
f_binary = wav_to_bytes("tests/vad_example.wav")
chunk_size = 960000
# chunk_size = len(f_binary)
print(f"f_binary: {len(f_binary)}, chunk_size: {chunk_size}, clip_num: {len(f_binary) // chunk_size}")
for i in range(0, len(f_binary), chunk_size):
binary_data = f_binary[i:i+chunk_size]
input_queue.put(binary_data)
# 等待VAD函数器结束
time.sleep(10)
vad_functor.stop()