diff --git a/README.md b/README.md new file mode 100644 index 0000000..e69de29 diff --git a/openwakeword/__init__.py b/openwakeword/__init__.py new file mode 100644 index 0000000..ba32899 --- /dev/null +++ b/openwakeword/__init__.py @@ -0,0 +1 @@ +from openwakeword.detect import Model \ No newline at end of file diff --git a/openwakeword/detect.py b/openwakeword/detect.py new file mode 100644 index 0000000..51afdf6 --- /dev/null +++ b/openwakeword/detect.py @@ -0,0 +1,51 @@ +# Copyright 2022 David Scripka. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Imports +import onnxruntime as ort +import numpy as np +import os +from openwakeword.utils import AudioFeatures +from typing import List + +class Model(): + def __init__(self, wakeword_model_paths: List[str], input_sizes: List[int], **kwargs): + # Initialize the ONNX models and store them + sessionOptions = ort.SessionOptions() + sessionOptions.inter_op_num_threads = 1 + sessionOptions.intra_op_num_threads = 1 + + self.models = {} + self.model_inputs = {} + for size, mdl_path in zip(input_sizes, wakeword_model_paths): + mdl_name = mdl_path.split(os.path.sep)[-1].strip(".onnx") + self.models[mdl_name] = ort.InferenceSession(mdl_path, sess_options=sessionOptions) + self.model_inputs[mdl_name] = size + + # Create AudioFeatures object + self.preprocessor = AudioFeatures(**kwargs) + + def predict(self, x): + """Predict with all of the wakeword models on the input audio frames""" + self.preprocessor(x) + predictions = {} + for mdl in self.models.keys(): + input_name = self.models[mdl].get_inputs()[0].name + predictions[mdl] = self.models[mdl].run( + None, + {input_name: self.preprocessor.get_features(self.model_inputs[mdl])} + )[0][0][0] + return predictions + + \ No newline at end of file diff --git a/openwakeword/resources/models/embedding_model.onnx b/openwakeword/resources/models/embedding_model.onnx new file mode 100644 index 0000000..8582948 --- /dev/null +++ b/openwakeword/resources/models/embedding_model.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ba754db3cd768a524c655ea90655ee5e6055a43b8dfd29366a11e93716ae9e51 +size 1328103 diff --git a/openwakeword/resources/models/hey_jane.onnx b/openwakeword/resources/models/hey_jane.onnx new file mode 100644 index 0000000..d05de2a --- /dev/null +++ b/openwakeword/resources/models/hey_jane.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7a76e2b151e2b6416b422c6c905b43b0dc20fca2774597c1a814cd467ea29966 +size 504084 diff --git a/openwakeword/resources/models/melspectrogram.onnx b/openwakeword/resources/models/melspectrogram.onnx new file mode 100644 index 0000000..d94407b --- /dev/null +++ b/openwakeword/resources/models/melspectrogram.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e942d7a12a2f10bc096b960fc06a31939400f64410f72d808acf5c394f09b863 +size 1087211 diff --git a/openwakeword/utils.py b/openwakeword/utils.py new file mode 100644 index 0000000..e932646 --- /dev/null +++ b/openwakeword/utils.py @@ -0,0 +1,102 @@ +# Copyright 2022 David Scripka. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Imports +import os +import onnxruntime as ort +import numpy as np +from collections import deque + +# Base class for computing audio features using Google's speech_embedding model (https://tfhub.dev/google/speech_embedding/1) +class AudioFeatures(): + def __init__(self, + melspec_onnx_model_path=os.path.join("openwakeword", "resources", "models", "melspectrogram.onnx"), + embedding_onnx_model_path=os.path.join("openwakeword", "resources", "models", "embedding_model.onnx"), + sr=16000, + ncpu=1 + ): + # Initialize the ONNX models + sessionOptions = ort.SessionOptions() + sessionOptions.inter_op_num_threads = ncpu + sessionOptions.intra_op_num_threads = ncpu + self.melspec_model = ort.InferenceSession(melspec_onnx_model_path, sess_options=sessionOptions) + self.embedding_model = ort.InferenceSession(embedding_onnx_model_path, sess_options=sessionOptions) + + # Create databuffers + self.raw_data_buffer = deque(maxlen=sr*10) + self.melspectrogram_buffer = np.zeros((0,32)) #n_frames x num_features + self.melspectrogram_max_len = 10*97 # 97 is the number of frames in 1 second of 16hz audio + self.feature_buffer = np.zeros((32,96)) + self.feature_buffer_max_len = 120 # ~10 seconds of feature buffer history + + def _get_melspectrogram(self, x): + """Function to compute the mel-spectrogram of the provided audio samples.""" + x = np.array(x)[None,] if isinstance(x, list) else x[None,] + x = x.astype(np.float32) if x.dtype!=np.float32 else x + outputs = self.melspec_model.run(None, {'input': x}) + spec = np.squeeze(outputs[0]) + spec = spec/10 + 2 # Arbitrary adjustment to make result closer to original Google speech_embedding model + + return spec + + def _get_embeddings(self, x, window_size=76, step_size = 8): + """Function to compute the embeddings of the provide audio samples.""" + spec = self._get_melspectrogram(x) + windows = [] + for i in range(0, spec.shape[0], 8): + window = spec[i:i+window_size] + if window.shape[0] == window_size: # truncate short windows + windows.append(window) + + batch = np.expand_dims(np.array(windows), axis=-1) + embedding = self.embedding_model.run(None, {'input_1': batch})[0].squeeze() + return embedding + + def _streaming_melspectrogram(self, x): + """Note! There seem to be some slight numerical issues depending on the underlying audio data + such that the streaming method is not exactly the same as when the melspectrogram of the entire + clip is calculated. It's unclear if this difference is significant and will impact model performance. + In particular padding with 0 or very small values seems to demonstrate the differences well. + """ + if len(x) < 400: + raise ValueError("The number of input frames must be at least 400 samples @ 16khz (25 ms)!") + self.raw_data_buffer.extend(x.tolist() if isinstance(x, np.ndarray) else x) + self.melspectrogram_buffer = np.vstack( + (self.melspectrogram_buffer, self._get_melspectrogram(list(self.raw_data_buffer)[-len(x)-160*3:])) + ) + + if self.melspectrogram_buffer.shape[0] > self.melspectrogram_max_len: + self.melspectrogram_buffer = self.melspectrogram_buffer[-self.melspectrogram_max_len:, :] + + def _streaming_features(self, x): + assert len(x) == 1280*3 + self._streaming_melspectrogram(x) + if self.melspectrogram_buffer.shape[0] < 76: + pass + for i in range(-3,0,1): + start = -76+8*i + end = start + 76 + x = self.melspectrogram_buffer[start:end].astype(np.float32)[None,:,:,None] + if x.shape[1] != 76: + continue + self.feature_buffer = np.vstack((self.feature_buffer, self.embedding_model.run(None, {'input_1': x})[0].squeeze())) + + if self.feature_buffer.shape[0] > self.feature_buffer_max_len: + self.feature_buffer = self.feature_buffer[-self.feature_buffer_max_len:, :] + + def get_features(self, x, n_feature_frames=16): + return self.feature_buffer[-n_feature_frames:, :][None,].astype(np.float32) + + def __call__(self, x): + self._streaming_features(x) \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..b0f0765 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,3 @@ +[build-system] +requires = ["setuptools>=42"] +build-backend = "setuptools.build_meta" diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..6b4ca1e --- /dev/null +++ b/setup.py @@ -0,0 +1,25 @@ +import setuptools + +with open("README.md", "r", encoding="utf-8") as fh: + long_description = fh.read() + +setuptools.setup( + name="openwakeword", + version="0.0.1", + author="David Scripka", + author_email="david.scripka@gmail.com", + description="An open-source audio wake word (or phrase) detection framework with a focus on accuracy and customizability", + long_description=long_description, + long_description_content_type="text/markdown", + url="https://github.com/pypa/sampleproject", + project_urls={ + "Bug Tracker": "https://github.com/pypa/sampleproject/issues", + }, + classifiers=[ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: Apache 2.0 License", + "Operating System :: OS Independent", + ], + packages=setuptools.find_packages(), + python_requires=">=3.6", +) \ No newline at end of file diff --git a/tests/data/hey_jane.wav b/tests/data/hey_jane.wav new file mode 100644 index 0000000..4daf992 Binary files /dev/null and b/tests/data/hey_jane.wav differ diff --git a/tests/test_models.py b/tests/test_models.py new file mode 100644 index 0000000..096af2c --- /dev/null +++ b/tests/test_models.py @@ -0,0 +1,38 @@ +# Copyright 2022 David Scripka. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Imports +import openwakeword +import scipy.io.wavfile +import pytest + +# Tests +class TestModels: + @pytest.fixture(scope="class") + def hey_jane_clip(self): + sr, dat = scipy.io.wavfile.read("tests/data/hey_jane.wav", "rb") + return dat + + def test_hey_jane(self, hey_jane_clip): + model = openwakeword.Model( + wakeword_model_paths=["openwakeword/resources/models/hey_jane.onnx"], + input_sizes=[16] + ) + + step_size = 1280*3 + predictions = [] + for i in range(0, hey_jane_clip.shape[0]-step_size, step_size): + predictions.append(model.predict(hey_jane_clip[i:i+step_size])["hey_jane"]) + + assert max(predictions) > 0.5