diff --git a/examples/detect_from_microphone.py b/examples/detect_from_microphone.py index eb102fb..9754d68 100644 --- a/examples/detect_from_microphone.py +++ b/examples/detect_from_microphone.py @@ -27,11 +27,8 @@ mic_stream = sounddevice.InputStream( dtype = np.int16, ) -# Load openwakeword model(s) -model_name = "alexa_v5" -model = Model( - wakeword_model_paths=[os.path.join("../", "openwakeword", "resources", "models", model_name + ".onnx")], -) +# Load pre-trained openwakeword models +owwModel = Model() # Run capture loop, checking for hotwords if __name__ == "__main__": @@ -39,20 +36,26 @@ if __name__ == "__main__": mic_stream.start() # Create a prediction buffer - prediction_buffer = [0]*30 while True: # Get audio audio, overflowed = mic_stream.read(1280) audio = audio.squeeze() # Feed to openWakeWord model - prediction = model.predict(audio) - prediction_buffer = prediction_buffer[1:] + [round(prediction[model_name], 2)] - - # Plot predictions in graph + prediction = owwModel.predict(audio) + + # Get predictions from prediction buffers and plot plt.cld() plt.clt() - plt.plot(prediction_buffer) + for mdl in owwModel.prediction_buffer.keys(): + # Plot scores in graph + scores = list(owwModel.prediction_buffer[mdl]) + plt.plot(scores) + + # Plot text showing name of model with scores >= 0.5 (default threshold) + if max(scores) >= 0.5: + plt.text(mdl, 15, 0.9, alignment="center", color = "blue", style="bold") + plt.ylim(0,1) plt.show() - plt.sleep(0.005) \ No newline at end of file + plt.sleep(0.005) diff --git a/openwakeword/__init__.py b/openwakeword/__init__.py index 904866d..227e1fc 100644 --- a/openwakeword/__init__.py +++ b/openwakeword/__init__.py @@ -1,3 +1,20 @@ +import os from openwakeword.model import Model __all__ = ['Model', ] + +models = { + "alexa": { + "model_path": os.path.join(os.path.dirname(os.path.abspath(__file__)), "resources/models/alexa_v5.onnx") + }, + "hey_mycroft": { + "model_path": os.path.join(os.path.dirname(os.path.abspath(__file__)), "resources/models/hey_mycroft_v1.onnx") + }, + "timer": { + "model_path": os.path.join(os.path.dirname(os.path.abspath(__file__)), "resources/models/timer_v1.onnx") + } +} + + +def get_pretrained_model_paths(): + return [models[i]["model_path"] for i in models.keys()] diff --git a/openwakeword/data.py b/openwakeword/data.py index 327bcf5..c80e2f9 100644 --- a/openwakeword/data.py +++ b/openwakeword/data.py @@ -483,6 +483,7 @@ class mmap_batch_generator: return np.vstack(X), np.array(y) + # Function to remove empty rows from the end of a mmap array def trim_mmap(mmap_path): """ @@ -502,11 +503,11 @@ def trim_mmap(mmap_path): i -= 1 N_new = mmap_file1.shape[0] + i + 1 - + # Create new mmap_file and copy over data in batches output_file2 = mmap_path.strip(".npy") + "2.npy" mmap_file2 = open_memmap(output_file2, mode='w+', dtype=np.float32, - shape=(N_new, mmap_file1.shape[1], mmap_file1.shape[2])) + shape=(N_new, mmap_file1.shape[1], mmap_file1.shape[2])) for i in tqdm(range(0, mmap_file1.shape[0], 1024), total=mmap_file1.shape[0]//1024): if i + 1024 > N_new: @@ -520,4 +521,4 @@ def trim_mmap(mmap_path): os.remove(mmap_path) # Rename new mmap file to match original - os.rename(output_file2, mmap_path) \ No newline at end of file + os.rename(output_file2, mmap_path) diff --git a/openwakeword/model.py b/openwakeword/model.py index 1eae2c7..8ea8a08 100644 --- a/openwakeword/model.py +++ b/openwakeword/model.py @@ -15,9 +15,9 @@ # Imports import numpy as np import onnxruntime as ort +import openwakeword from openwakeword.utils import AudioFeatures -import statistics import wave import os import json @@ -34,12 +34,13 @@ class Model(): The main model class for openWakeWord. Creates a model object with the shared audio pre-processer and for arbitrarily many custom wake word/wake phrase models. """ - def __init__(self, wakeword_model_paths: List[str], **kwargs): + def __init__(self, wakeword_model_paths: List[str] = [], **kwargs): """ Initialize the openWakeWord model object. Args: - wakeword_model_paths (List[str]): A list of paths of ONNX models to load into the openWakeWord model object + wakeword_model_paths (List[str]): A list of paths of ONNX models to load into the openWakeWord model object. + If not provided, will load all of the pre-trained models. """ # Initialize the ONNX models and store them @@ -47,14 +48,20 @@ class Model(): sessionOptions.inter_op_num_threads = 1 sessionOptions.intra_op_num_threads = 1 + # Get model paths for pre-trained models if user doesn't provide models to load + if wakeword_model_paths == []: + wakeword_model_paths = openwakeword.get_pretrained_model_paths() + wakeword_model_names = list(openwakeword.models.keys()) + else: + wakeword_model_names = [os.path.basename(i.strip(".onnx")) for i in wakeword_model_paths] + # Create attributes to store models and metadat self.models = {} self.model_inputs = {} self.model_outputs = {} self.class_mapping = {} self.model_input_names = {} - for mdl_path in wakeword_model_paths: - mdl_name = mdl_path.split(os.path.sep)[-1].strip(".onnx") + for mdl_path, mdl_name in zip(wakeword_model_paths, wakeword_model_names): self.models[mdl_name] = ort.InferenceSession(mdl_path, sess_options=sessionOptions, providers=["CPUExecutionProvider"]) self.model_inputs[mdl_name] = self.models[mdl_name].get_inputs()[0].shape[1] @@ -96,10 +103,11 @@ class Model(): x (Union[ndarray, List]): The input audio data to predict on with the models. Must be 1280 samples of 16khz, 16-bit audio data. patience (dict): How many consecutive frames (of 1280 samples or 80 ms) above the threshold that must - be observed before the current frame will be returned as non-zero. + be observed before the current frame will be returned as non-zero. Must be provided as an a dictionary where the keys are the - model names and the values are the number of frames. Can reduce false-positive detections at the - cost of a lower true-positive rate. By default, this behavior is disabled. + model names and the values are the number of frames. Can reduce false-positive + detections at the cost of a lower true-positive rate. + By default, this behavior is disabled. threshold (dict): The threshold values to use when the `patience` behavior is enabled. Must be provided as an a dictionary where the keys are the model names and the values are the thresholds. diff --git a/tests/test_models.py b/tests/test_models.py index c3bbe74..9bacfa6 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -33,30 +33,24 @@ import numpy as np from pathlib import Path import collections -# Define models and corresponding files for testing -test_dict = { - "hey_mycroft_v1": ["hey_mycroft_v1_test.wav"], - "alexa_v5": ["alexa_v5_test.wav"] -} - # Tests class TestModels: def test_models(self): # Load model - models = [str(i) for i in Path( - os.path.join("openwakeword", "resources", "models") - ).glob("**/*.onnx") - if "embedding" not in str(i) and "melspec" not in str(i)] - owwModel = openwakeword.Model( - wakeword_model_paths=models, - ) + owwModel = openwakeword.Model() + + # Get clips for each model (assumes that test clips will have the model name in the filename) + test_dict = {} + for mdl_name in owwModel.models.keys(): + all_clips = [str(i) for i in Path(os.path.join("tests", "data")).glob("*.wav")] + test_dict[mdl_name] = [i for i in all_clips if mdl_name in i] # Predict for model, clips in test_dict.items(): for clip in clips: # Get predictions for reach frame in the clip - predictions = owwModel.predict_clip(os.path.join("tests", "data", clip)) + predictions = owwModel.predict_clip(clip) owwModel.reset() # reset after each clip to ensure independent results # Make predictions dictionary flatter @@ -70,35 +64,6 @@ class TestModels: else: assert max(predictions_flat[key]) < 0.5 - def test_models_with_median_smooth(self): - # Load models - models = [str(i) for i in Path( - os.path.join("openwakeword", "resources", "models") - ).glob("**/*.onnx") - if "embedding" not in str(i) and "melspec" not in str(i)] - owwModel = openwakeword.Model( - wakeword_model_paths=models, - ) - - # Predict with median smooth - for model, clips in test_dict.items(): - for clip in clips: - # Get predictions for reach frame in the clip - predictions = owwModel.predict_clip(os.path.join("tests", "data", clip), median_smooth=True) - owwModel.reset() # reset after each clip to ensure independent results - - # Make predictions dictionary flatter - predictions_flat = collections.defaultdict(list) - [predictions_flat[key].append(i[key]) for i in predictions for key in i.keys()] - - # Check scores against default threshold (0.5) - for key in predictions_flat.keys(): - if key in clip: - assert max(predictions_flat[key]) >= 0.5 - else: - print(key, clip) - assert max(predictions_flat[key]) < 0.5 - def test_models_with_timing(self): models = [str(i) for i in Path( os.path.join("openwakeword", "resources", "models")