diff --git a/files/en_num.wav b/files/en_num.wav new file mode 100644 index 0000000..906ff26 Binary files /dev/null and b/files/en_num.wav differ diff --git a/files/number_detector.jit b/files/number_detector.jit new file mode 100644 index 0000000..3e9c9c2 Binary files /dev/null and b/files/number_detector.jit differ diff --git a/files/number_detector.onnx b/files/number_detector.onnx new file mode 100644 index 0000000..2a7e5e4 Binary files /dev/null and b/files/number_detector.onnx differ diff --git a/files/ru_num.wav b/files/ru_num.wav new file mode 100644 index 0000000..2f90971 Binary files /dev/null and b/files/ru_num.wav differ diff --git a/hubconf.py b/hubconf.py index 16756cb..c0f16a4 100644 --- a/hubconf.py +++ b/hubconf.py @@ -2,11 +2,13 @@ dependencies = ['torch', 'torchaudio'] import torch from utils import (init_jit_model, get_speech_ts, + get_number_ts, save_audio, read_audio, state_generator, single_audio_stream, - collect_speeches) + collect_chunks, + drop_chunks) def silero_vad(**kwargs): @@ -21,6 +23,22 @@ def silero_vad(**kwargs): read_audio, state_generator, single_audio_stream, - collect_speeches) + collect_chunks) + + return model, utils + + +def silero_number_detector(**kwargs): + """Silero Number Detector and Language Classifier + 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/number_detector.jit') + utils = (get_number_ts, + save_audio, + read_audio, + collect_chunks, + drop_chunks) return model, utils diff --git a/silero-vad.ipynb b/silero-vad.ipynb old mode 100644 new mode 100755 index 818b8e6..e90e9d0 --- a/silero-vad.ipynb +++ b/silero-vad.ipynb @@ -6,16 +6,26 @@ "id": "sVNOuHQQjsrp" }, "source": [ - "# PyTorch Example" + "# PyTorch Examples" ] }, { "cell_type": "markdown", "metadata": { - "id": "9ZTzCtc5kYVg" + "heading_collapsed": true }, "source": [ - "## Install Dependencies" + "## VAD" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "heading_collapsed": true, + "hidden": true + }, + "source": [ + "### Install Dependencies" ] }, { @@ -23,12 +33,10 @@ "execution_count": null, "metadata": { "ExecuteTime": { - "end_time": "2020-12-15T14:00:15.701867Z", - "start_time": "2020-12-15T14:00:09.512876Z" + "end_time": "2020-12-30T17:35:43.397137Z", + "start_time": "2020-12-30T17:33:10.962078Z" }, - "cellView": "form", - "collapsed": true, - "id": "rllMjjsekbjt" + "hidden": true }, "outputs": [], "source": [ @@ -53,7 +61,7 @@ " read_audio,\n", " state_generator,\n", " single_audio_stream,\n", - " collect_speeches) = utils\n", + " collect_chunks) = utils\n", "\n", "files_dir = torch.hub.get_dir() + '/snakers4_silero-vad_master/files'" ] @@ -61,10 +69,12 @@ { "cell_type": "markdown", "metadata": { + "heading_collapsed": true, + "hidden": true, "id": "fXbbaUO3jsrw" }, "source": [ - "## Full Audio" + "### Full Audio" ] }, { @@ -72,9 +82,10 @@ "execution_count": null, "metadata": { "ExecuteTime": { - "end_time": "2020-12-15T13:09:56.879818Z", - "start_time": "2020-12-15T13:09:56.864765Z" + "end_time": "2020-12-30T17:35:44.362860Z", + "start_time": "2020-12-30T17:35:43.398441Z" }, + "hidden": true, "id": "aI_eydBPjsrx" }, "outputs": [], @@ -91,26 +102,29 @@ "execution_count": null, "metadata": { "ExecuteTime": { - "end_time": "2020-12-15T13:09:58.941063Z", - "start_time": "2020-12-15T13:09:58.887006Z" + "end_time": "2020-12-30T17:35:44.419280Z", + "start_time": "2020-12-30T17:35:44.364175Z" }, + "hidden": true, "id": "OuEobLchjsry" }, "outputs": [], "source": [ "# merge all speech chunks to one audio\n", "save_audio('only_speech.wav',\n", - " collect_speeches(speech_timestamps, wav), 16000) \n", + " collect_chunks(speech_timestamps, wav), 16000) \n", "Audio('only_speech.wav')" ] }, { "cell_type": "markdown", "metadata": { + "heading_collapsed": true, + "hidden": true, "id": "iDKQbVr8jsry" }, "source": [ - "## Single Audio Stream" + "### Single Audio Stream" ] }, { @@ -121,6 +135,7 @@ "end_time": "2020-12-15T13:09:59.199321Z", "start_time": "2020-12-15T13:09:59.196823Z" }, + "hidden": true, "id": "q-lql_2Wjsry" }, "outputs": [], @@ -135,10 +150,12 @@ { "cell_type": "markdown", "metadata": { + "heading_collapsed": true, + "hidden": true, "id": "KBDVybJCjsrz" }, "source": [ - "## Multiple Audio Streams" + "### Multiple Audio Streams" ] }, { @@ -149,6 +166,7 @@ "end_time": "2020-12-15T13:10:03.590358Z", "start_time": "2020-12-15T13:10:03.587071Z" }, + "hidden": true, "id": "BK4tGfWgjsrz" }, "outputs": [], @@ -165,6 +183,7 @@ "end_time": "2020-12-15T13:10:15.762491Z", "start_time": "2020-12-15T13:10:03.591388Z" }, + "hidden": true, "id": "v1l8sam1jsrz" }, "outputs": [], @@ -174,6 +193,125 @@ " pprint(batch)" ] }, + { + "cell_type": "markdown", + "metadata": { + "heading_collapsed": true + }, + "source": [ + "## Number detector" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "heading_collapsed": true, + "hidden": true + }, + "source": [ + "### Install Dependencies" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "hidden": true + }, + "outputs": [], + "source": [ + "#@title Install and Import Dependencies\n", + "\n", + "# this assumes that you have a relevant version of PyTorch installed\n", + "!pip install -q torchaudio soundfile\n", + "\n", + "import glob\n", + "import torch\n", + "torch.set_num_threads(1)\n", + "\n", + "from IPython.display import Audio\n", + "from pprint import pprint\n", + "\n", + "model, utils = torch.hub.load(repo_or_dir='snakers4/silero-vad',\n", + " model='silero_number_detector',\n", + " force_reload=True)\n", + "\n", + "(get_number_ts,\n", + " save_audio,\n", + " read_audio,\n", + " collect_chunks,\n", + " drop_chunks) = utils\n", + "\n", + "files_dir = torch.hub.get_dir() + '/snakers4_silero-vad_master/files'" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "heading_collapsed": true, + "hidden": true + }, + "source": [ + "### Full audio" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "hidden": true + }, + "outputs": [], + "source": [ + "wav = read_audio(f'{files_dir}/en_num.wav')\n", + "# get number timestamps from full audio file\n", + "number_timestamps = get_number_ts(wav, model)\n", + "pprint(number_timestamps)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "hidden": true + }, + "outputs": [], + "source": [ + "sample_rate = 16000\n", + "# convert ms in timestamps to samples\n", + "for timestamp in number_timestamps:\n", + " timestamp['start'] = int(timestamp['start'] * sample_rate)\n", + " timestamp['end'] = int(timestamp['end'] * sample_rate)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "hidden": true + }, + "outputs": [], + "source": [ + "# merge all number chunks to one audio\n", + "save_audio('only_numbers.wav',\n", + " collect_chunks(number_timestamps, wav), sample_rate) \n", + "Audio('only_numbers.wav')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "hidden": true + }, + "outputs": [], + "source": [ + "# drop all number chunks from audio\n", + "save_audio('no_numbers.wav',\n", + " drop_chunks(number_timestamps, wav), sample_rate) \n", + "Audio('no_numbers.wav')" + ] + }, { "cell_type": "markdown", "metadata": { @@ -186,10 +324,21 @@ { "cell_type": "markdown", "metadata": { + "heading_collapsed": true + }, + "source": [ + "## VAD" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "heading_collapsed": true, + "hidden": true, "id": "bL4kn4KJrlyL" }, "source": [ - "## Install Dependencies" + "### Install Dependencies" ] }, { @@ -197,6 +346,7 @@ "execution_count": null, "metadata": { "cellView": "form", + "hidden": true, "id": "Q4QIfSpprnkI" }, "outputs": [], @@ -239,10 +389,12 @@ { "cell_type": "markdown", "metadata": { + "heading_collapsed": true, + "hidden": true, "id": "5JHErdB7jsr0" }, "source": [ - "## Full Audio" + "### Full Audio" ] }, { @@ -253,6 +405,7 @@ "end_time": "2020-12-15T13:09:06.643812Z", "start_time": "2020-12-15T13:09:06.473386Z" }, + "hidden": true, "id": "krnGoA6Kjsr0" }, "outputs": [], @@ -273,22 +426,25 @@ "end_time": "2020-12-15T13:09:08.862421Z", "start_time": "2020-12-15T13:09:08.820014Z" }, + "hidden": true, "id": "B176Lzfnjsr1" }, "outputs": [], "source": [ "# merge all speech chunks to one audio\n", - "save_audio('only_speech.wav', collect_speeches(speech_timestamps, wav), 16000)\n", + "save_audio('only_speech.wav', collect_chunks(speech_timestamps, wav), 16000)\n", "Audio('only_speech.wav')" ] }, { "cell_type": "markdown", "metadata": { + "heading_collapsed": true, + "hidden": true, "id": "Rio9W50gjsr1" }, "source": [ - "## Single Audio Stream" + "### Single Audio Stream" ] }, { @@ -299,6 +455,7 @@ "end_time": "2020-12-15T13:09:09.606031Z", "start_time": "2020-12-15T13:09:09.504239Z" }, + "hidden": true, "id": "IPkl8Yy1jsr1" }, "outputs": [], @@ -315,6 +472,7 @@ "end_time": "2020-12-15T13:09:11.453171Z", "start_time": "2020-12-15T13:09:09.633435Z" }, + "hidden": true, "id": "NC6Jim0hjsr1" }, "outputs": [], @@ -327,10 +485,12 @@ { "cell_type": "markdown", "metadata": { + "heading_collapsed": true, + "hidden": true, "id": "WNZ42u0ajsr1" }, "source": [ - "## Multiple Audio Streams" + "### Multiple Audio Streams" ] }, { @@ -341,6 +501,7 @@ "end_time": "2020-12-15T13:09:11.540423Z", "start_time": "2020-12-15T13:09:11.455706Z" }, + "hidden": true, "id": "XjhGQGppjsr1" }, "outputs": [], @@ -358,6 +519,7 @@ "end_time": "2020-12-15T13:09:19.565434Z", "start_time": "2020-12-15T13:09:11.552097Z" }, + "hidden": true, "id": "QI7-arlqjsr2" }, "outputs": [], @@ -366,6 +528,154 @@ " if batch:\n", " pprint(batch)" ] + }, + { + "cell_type": "markdown", + "metadata": { + "heading_collapsed": true + }, + "source": [ + "## Number detector" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "heading_collapsed": true, + "hidden": true, + "id": "bL4kn4KJrlyL" + }, + "source": [ + "### Install Dependencies" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "ExecuteTime": { + "end_time": "2020-12-30T17:25:19.107534Z", + "start_time": "2020-12-30T17:24:51.853293Z" + }, + "cellView": "form", + "hidden": true, + "id": "Q4QIfSpprnkI" + }, + "outputs": [], + "source": [ + "#@title Install and Import Dependencies\n", + "\n", + "# this assumes that you have a relevant version of PyTorch installed\n", + "!pip install -q torchaudio soundfile onnxruntime\n", + "\n", + "import glob\n", + "import torch\n", + "import onnxruntime\n", + "from pprint import pprint\n", + "\n", + "from IPython.display import Audio\n", + "\n", + "_, utils = torch.hub.load(repo_or_dir='snakers4/silero-vad',\n", + " model='silero_number_detector',\n", + " force_reload=True)\n", + "\n", + "(get_number_ts,\n", + " save_audio,\n", + " read_audio,\n", + " collect_chunks,\n", + " drop_chunks) = utils\n", + "\n", + "files_dir = torch.hub.get_dir() + '/snakers4_silero-vad_master/files'\n", + "\n", + "def init_onnx_model(model_path: str):\n", + " return onnxruntime.InferenceSession(model_path)\n", + "\n", + "def validate_onnx(model, inputs):\n", + " with torch.no_grad():\n", + " ort_inputs = {'input': inputs.cpu().numpy()}\n", + " outs = model.run(None, ort_inputs)\n", + " outs = [torch.Tensor(x) for x in outs]\n", + " return outs" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "hidden": true, + "id": "5JHErdB7jsr0" + }, + "source": [ + "### Full Audio" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "ExecuteTime": { + "end_time": "2020-12-15T13:09:06.643812Z", + "start_time": "2020-12-15T13:09:06.473386Z" + }, + "hidden": true, + "id": "krnGoA6Kjsr0" + }, + "outputs": [], + "source": [ + "model = init_onnx_model(f'{files_dir}/number_detector.onnx')\n", + "wav = read_audio(f'{files_dir}/en_num.wav')\n", + "\n", + "# get number timestamps from full audio file\n", + "number_timestamps = get_number_ts(wav, model, run_function=validate_onnx)\n", + "pprint(number_timestamps)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "hidden": true + }, + "outputs": [], + "source": [ + "sample_rate = 16000\n", + "# convert ms in timestamps to samples\n", + "for timestamp in number_timestamps:\n", + " timestamp['start'] = int(timestamp['start'] * sample_rate)\n", + " timestamp['end'] = int(timestamp['end'] * sample_rate)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "ExecuteTime": { + "end_time": "2020-12-15T13:09:08.862421Z", + "start_time": "2020-12-15T13:09:08.820014Z" + }, + "hidden": true, + "id": "B176Lzfnjsr1" + }, + "outputs": [], + "source": [ + "# merge all number chunks to one audio\n", + "save_audio('only_numbers.wav',\n", + " collect_chunks(number_timestamps, wav), 16000) \n", + "Audio('only_numbers.wav')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "hidden": true + }, + "outputs": [], + "source": [ + "# drop all number chunks from audio\n", + "save_audio('no_numbers.wav',\n", + " drop_chunks(number_timestamps, wav), 16000) \n", + "Audio('no_numbers.wav')" + ] } ], "metadata": { diff --git a/utils.py b/utils.py index 5c6bafe..597f4b5 100644 --- a/utils.py +++ b/utils.py @@ -105,6 +105,39 @@ def get_speech_ts(wav: torch.Tensor, return speeches +def get_number_ts(wav: torch.Tensor, + model, + model_stride=8, + hop_length=160, + sample_rate=16000, + run_function=validate): + wav = torch.unsqueeze(wav, dim=0) + perframe_logits = run_function(model, wav)[0] + perframe_preds = torch.argmax(torch.softmax(perframe_logits, dim=1), dim=1).squeeze() # (1, num_frames_strided) + extended_preds = [] + for i in perframe_preds: + extended_preds.extend([i.item()] * model_stride) + # len(extended_preds) is *num_frames_real*; for each frame of audio we know if it has a number in it. + triggered = False + timings = [] + cur_timing = {} + for i, pred in enumerate(extended_preds): + if pred == 1: + if not triggered: + cur_timing['start'] = (i * hop_length) / sample_rate + triggered = True + elif pred == 0: + if triggered: + cur_timing['end'] = (i * hop_length) / sample_rate + timings.append(cur_timing) + cur_timing = {} + triggered = False + if cur_timing: + cur_timing['end'] = len(wav) / sample_rate + timings.append(cur_timing) + return timings + + class VADiterator: def __init__(self, trig_sum: float = 0.26, @@ -252,9 +285,19 @@ def single_audio_stream(model, yield states -def collect_speeches(tss: List[dict], - wav: torch.Tensor): - speech_chunks = [] +def collect_chunks(tss: List[dict], + wav: torch.Tensor): + chunks = [] for i in tss: - speech_chunks.append(wav[i['start']: i['end']]) - return torch.cat(speech_chunks) + chunks.append(wav[i['start']: i['end']]) + return torch.cat(chunks) + + +def drop_chunks(tss: List[dict], + wav: torch.Tensor): + chunks = [] + cur_start = 0 + for i in tss: + chunks.append((wav[cur_start: i['start']])) + cur_start = i['end'] + return torch.cat(chunks)