mirror of
https://github.com/snakers4/silero-vad.git
synced 2026-02-04 17:39:22 +08:00
Use the path to hubconf.py to find models
While `torch.hub.load` uses predictable names when downloading from github, the name referenced only worked when using the `master` branch of the repo. Using a ref like `snakers4/silero_vad:v4.0` results in a different directory name and the model files not being found. Since the model files are stored in a path relative to hubconf.py, use the path of hubconf.py to find the models instead of assuming the directory `torch.hub.load` will extract the files to.
This commit is contained in:
14
hubconf.py
14
hubconf.py
@@ -1,6 +1,7 @@
|
||||
dependencies = ['torch', 'torchaudio']
|
||||
import torch
|
||||
import json
|
||||
import os
|
||||
from utils_vad import (init_jit_model,
|
||||
get_speech_timestamps,
|
||||
get_number_ts,
|
||||
@@ -31,11 +32,11 @@ def silero_vad(onnx=False, force_onnx_cpu=False):
|
||||
if versiontuple(installed_version) < versiontuple(supported_version):
|
||||
raise Exception(f'Please install torch {supported_version} or greater ({installed_version} installed)')
|
||||
|
||||
hub_dir = torch.hub.get_dir()
|
||||
model_dir = os.path.join(os.path.dirname(__file__), 'files')
|
||||
if onnx:
|
||||
model = OnnxWrapper(f'{hub_dir}/snakers4_silero-vad_master/files/silero_vad.onnx', force_onnx_cpu)
|
||||
model = OnnxWrapper(os.path.join(model_dir, 'silero_vad.onnx'), force_onnx_cpu)
|
||||
else:
|
||||
model = init_jit_model(model_path=f'{hub_dir}/snakers4_silero-vad_master/files/silero_vad.jit')
|
||||
model = init_jit_model(os.path.join(model_dir, 'silero_vad.jit'))
|
||||
utils = (get_speech_timestamps,
|
||||
save_audio,
|
||||
read_audio,
|
||||
@@ -85,18 +86,17 @@ def silero_lang_detector_95(onnx=False, force_onnx_cpu=False):
|
||||
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()
|
||||
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, force_onnx_cpu)
|
||||
|
||||
with open(f'{hub_dir}/snakers4_silero-vad_master/files/lang_dict_95.json', 'r') as f:
|
||||
model_dir = os.path.join(os.path.dirname(__file__), 'files')
|
||||
with open(os.path.join(model_dir, 'lang_dict_95.json'), 'r') as f:
|
||||
lang_dict = json.load(f)
|
||||
|
||||
with open(f'{hub_dir}/snakers4_silero-vad_master/files/lang_group_dict_95.json', 'r') as f:
|
||||
with open(os.path.join(model_dir, 'lang_group_dict_95.json'), 'r') as f:
|
||||
lang_group_dict = json.load(f)
|
||||
|
||||
utils = (get_language_and_group, read_audio)
|
||||
|
||||
Reference in New Issue
Block a user