Initial commit, basic model implemented and first test passing

openwakeword/__init__.py

@@ -0,0 +1 @@
+from openwakeword.detect import Model

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
+
+        

openwakeword/resources/models/embedding_model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ba754db3cd768a524c655ea90655ee5e6055a43b8dfd29366a11e93716ae9e51
+size 1328103

openwakeword/resources/models/hey_jane.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7a76e2b151e2b6416b422c6c905b43b0dc20fca2774597c1a814cd467ea29966
+size 504084

openwakeword/resources/models/melspectrogram.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e942d7a12a2f10bc096b960fc06a31939400f64410f72d808acf5c394f09b863
+size 1087211

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)

pyproject.toml

@@ -0,0 +1,3 @@
+[build-system]
+requires = ["setuptools>=42"]
+build-backend = "setuptools.build_meta"

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",
+)

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