add onnx vad

This commit is contained in:
adamnsandle
2021-12-17 14:48:32 +00:00
parent 0feae6cbbe
commit 74f759c8f8
10 changed files with 185 additions and 377 deletions

View File

@@ -1,7 +1,6 @@
dependencies = ['torch', 'torchaudio']
import torch
import json
from utils_vad import (init_jit_model,
get_speech_timestamps,
get_number_ts,
@@ -12,16 +11,20 @@ from utils_vad import (init_jit_model,
VADIterator,
collect_chunks,
drop_chunks,
donwload_onnx_model)
Validator,
OnnxWrapper)
def silero_vad(**kwargs):
def silero_vad(onnx=False):
"""Silero Voice Activity Detector
Returns a model with a set of utils
Please see https://github.com/snakers4/silero-vad for usage examples
"""
hub_dir = torch.hub.get_dir()
model = init_jit_model(model_path=f'{hub_dir}/snakers4_silero-vad_master/files/silero_vad.jit')
if onnx:
model = OnnxWrapper(f'{hub_dir}/snakers4_silero-vad_master/files/silero_vad.onnx')
else:
model = init_jit_model(model_path=f'{hub_dir}/snakers4_silero-vad_master/files/silero_vad.jit')
utils = (get_speech_timestamps,
save_audio,
read_audio,
@@ -31,46 +34,53 @@ def silero_vad(**kwargs):
return model, utils
def silero_number_detector(**kwargs):
def silero_number_detector(onnx=False):
"""Silero Number Detector
Returns a model with a set of utils
Please see https://github.com/snakers4/silero-vad for usage examples
"""
torch.hub.download_url_to_file('https://models.silero.ai/vad_models/number_detector.jit', 'number_detector.jit')
model = init_jit_model(model_path='number_detector.jit')
if onnx:
url = 'https://models.silero.ai/vad_models/number_detector.onnx'
else:
url = 'https://models.silero.ai/vad_models/number_detector.jit'
model = Validator(url)
utils = (get_number_ts,
save_audio,
read_audio,
collect_chunks,
drop_chunks,
donwload_onnx_model)
drop_chunks)
return model, utils
def silero_lang_detector(**kwargs):
def silero_lang_detector(onnx=False):
"""Silero Language Classifier
Returns a model with a set of utils
Please see https://github.com/snakers4/silero-vad for usage examples
"""
torch.hub.download_url_to_file('https://models.silero.ai/vad_models/number_detector.jit', 'number_detector.jit')
model = init_jit_model(model_path='number_detector.jit')
if onnx:
url = 'https://models.silero.ai/vad_models/number_detector.onnx'
else:
url = 'https://models.silero.ai/vad_models/number_detector.jit'
model = Validator(url)
utils = (get_language,
read_audio,
donwload_onnx_model)
read_audio)
return model, utils
def silero_lang_detector_95(**kwargs):
def silero_lang_detector_95(onnx=False):
"""Silero Language Classifier (95 languages)
Returns a model with a set of utils
Please see https://github.com/snakers4/silero-vad for usage examples
"""
hub_dir = torch.hub.get_dir()
torch.hub.download_url_to_file('https://models.silero.ai/vad_models/lang_classifier_95.jit', 'lang_classifier_95.jit')
model = init_jit_model(model_path='lang_classifier_95.jit')
if onnx:
url = 'https://models.silero.ai/vad_models/lang_classifier_95.onnx'
else:
url = 'https://models.silero.ai/vad_models/lang_classifier_95.jit'
model = Validator(url)
with open(f'{hub_dir}/snakers4_silero-vad_master/files/lang_dict_95.json', 'r') as f:
lang_dict = json.load(f)
@@ -78,6 +88,6 @@ def silero_lang_detector_95(**kwargs):
with open(f'{hub_dir}/snakers4_silero-vad_master/files/lang_group_dict_95.json', 'r') as f:
lang_group_dict = json.load(f)
utils = (get_language_and_group, read_audio, donwload_onnx_model)
utils = (get_language_and_group, read_audio)
return model, lang_dict, lang_group_dict, utils