diff --git a/openwakeword/data.py b/openwakeword/data.py index 1918041..327bcf5 100644 --- a/openwakeword/data.py +++ b/openwakeword/data.py @@ -21,6 +21,7 @@ from tqdm import tqdm from typing import List import numpy as np import torch +from numpy.lib.format import open_memmap from speechbrain.dataio.dataio import read_audio from speechbrain.processing.signal_processing import reverberate import torchaudio @@ -335,6 +336,9 @@ def mix_clips_batch( # Convert to 16-bit PCM audio mixed_clips_batch = (mixed_clips_batch.numpy()*32767).astype(np.int16) + # Remove any clips that are silent (happens rarely when mixing/reverberating) + mixed_clips_batch = mixed_clips_batch[np.where(mixed_clips_batch.max(axis=1) != 0)[0]] + yield mixed_clips_batch @@ -415,7 +419,7 @@ class mmap_batch_generator: self.data_transform_funcs = data_transform_funcs self.label_transform_funcs = label_transform_funcs - # Get array mmaps and store their shapes + # Get array mmaps and store their shapes (but load files < 1 GB total size into memory) self.data = {label: np.load(fl, mmap_mode='r') for label, fl in data_files.items()} self.data_counter = {label: 0 for label in data_files.keys()} self.original_shapes = {label: self.data[label].shape for label in self.data.keys()} @@ -456,7 +460,7 @@ class mmap_batch_generator: # Restart at zeroth index if an array reaches the end if self.data_counter[label] >= self.shapes[label][0]: self.data_counter[label] = 0 - self.data[label] = np.load(self.data_files[label], mmap_mode='r') + # self.data[label] = np.load(self.data_files[label], mmap_mode='r') # Get data from mmaped file x = self.data[label][self.data_counter[label]:self.data_counter[label]+n] @@ -478,3 +482,42 @@ class mmap_batch_generator: y.extend(y_batch) return np.vstack(X), np.array(y) + +# Function to remove empty rows from the end of a mmap array +def trim_mmap(mmap_path): + """ + Trims blank rows from the end of a mmaped numpy array by creates new mmap array without the blank rows. + Note that a copy is created and disk usage will briefly double as the function runs. + + Args: + mmap_path (str): The path to mmap array file to trim + + Returns: + None + """ + # Identify the last full row in the mmaped file + mmap_file1 = np.load(mmap_path, mmap_mode='r') + i = -1 + while np.all(mmap_file1[i, :, :] == 0): + 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])) + + for i in tqdm(range(0, mmap_file1.shape[0], 1024), total=mmap_file1.shape[0]//1024): + if i + 1024 > N_new: + mmap_file2[i:N_new] = mmap_file1[i:N_new].copy() + mmap_file2.flush() + else: + mmap_file2[i:i+1024] = mmap_file1[i:i+1024].copy() + mmap_file2.flush() + + # Remove old mmaped file + os.remove(mmap_path) + + # Rename new mmap file to match original + os.rename(output_file2, mmap_path) \ No newline at end of file diff --git a/openwakeword/model.py b/openwakeword/model.py index 4cacbd8..1eae2c7 100644 --- a/openwakeword/model.py +++ b/openwakeword/model.py @@ -20,6 +20,7 @@ from openwakeword.utils import AudioFeatures import statistics import wave import os +import json from collections import deque, defaultdict from functools import partial import time @@ -49,29 +50,59 @@ class Model(): # 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") 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] + self.model_outputs[mdl_name] = self.models[mdl_name].get_outputs()[0].shape[1] + output_name = self.models[mdl_name].get_outputs()[0].name + if "{" in output_name: + self.class_mapping[mdl_name] = json.loads(output_name) + else: + self.class_mapping[mdl_name] = mdl_name self.model_input_names[mdl_name] = self.models[mdl_name].get_inputs()[0].name # Create buffer to store frame predictios - self.prediction_buffer: DefaultDict[str, deque] = defaultdict(partial(deque, maxlen=5)) + self.prediction_buffer: DefaultDict[str, deque] = defaultdict(partial(deque, maxlen=30)) # Create AudioFeatures object self.preprocessor = AudioFeatures(**kwargs) - def predict(self, x: Union[np.ndarray, List], median_smooth: bool = False, timing: bool = False): + def get_parent_model_from_label(self, label): + """Gets the parent model associated with a given prediction label""" + parent_model = "" + for mdl in self.class_mapping.keys(): + if isinstance(self.class_mapping[mdl], dict): + if label in self.class_mapping[mdl].values(): + parent_model = mdl + else: + if label == self.class_mapping[mdl]: + parent_model = self.class_mapping[mdl] + + return parent_model + + def reset(self): + """Reset the prediction buffer""" + self.prediction_buffer = defaultdict(partial(deque, maxlen=30)) + + def predict(self, x: Union[np.ndarray, List], patience: dict = {}, threshold: dict = {}, timing: bool = False): """Predict with all of the wakeword models on the input audio frames Args: x (Union[ndarray, List]): The input audio data to predict on with the models. Must be 1280 samples of 16khz, 16-bit audio data. - median_smooth (bool): Whether to apply a running median smooth of the last three predictions - before returning a score. Can reduce false-positive productions at the - cost of a lower true-positive rate. + 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. + 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. + 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. timing (bool): Whether to print timing information of the models. Can be useful to debug and assess how efficiently models are running the current hardware. @@ -100,36 +131,53 @@ class Model(): model_start = time.time() # Run model - predictions[mdl] = self.models[mdl].run( + prediction = self.models[mdl].run( None, {input_name: self.preprocessor.get_features(self.model_inputs[mdl])} - )[0][0][0] + ) + if self.model_outputs[mdl] == 1: + predictions[mdl] = prediction[0][0][0] + else: + for int_label, cls in self.class_mapping[mdl].items(): + predictions[cls] = prediction[0][0][int(int_label)] - # Update prediction buffer - self.prediction_buffer[mdl].append(predictions[mdl]) - - # (Optionally) Smooth model predictions with simple median calculate of last three predictions - if median_smooth: - predictions[mdl] = statistics.median(list(self.prediction_buffer[mdl])[-3:]) + # Update prediction buffer, and zero predictions for first 5 frames during model initialization + for mdl in predictions.keys(): + if len(self.prediction_buffer[mdl]) < 5: + predictions[mdl] = 0.0 + self.prediction_buffer[mdl].append(predictions[mdl]) + # Get timing information if timing: model_end = time.time() timing_dict["models"][mdl] = model_end - model_start + # Update scores based on thresholds or patience arguments + if patience != {}: + if threshold == {}: + raise ValueError("Error! When using the `patience` argument, threshold " + "values must be provided via the `threshold` argument!") + for mdl in predictions.keys(): + parent_model = self.get_parent_model_from_label(mdl) + if parent_model in patience.keys(): + scores = np.array(self.prediction_buffer[mdl])[-patience[parent_model]:] + if (scores >= threshold[parent_model]).sum() < patience[parent_model]: + predictions[mdl] = 0.0 + if timing: pprint.PrettyPrinter().pprint(timing_dict) return predictions else: return predictions - def predict_clip(self, clip: str, padding: bool = True, **kwargs): + def predict_clip(self, clip: str, padding: int = 1, **kwargs): """Predict on an full audio clip, simulating streaming prediction. The input clip must bit a 16-bit, 16 khz, single-channel WAV file. Args: clip (str): The path to a 16-bit PCM, 16 khz, single-channel WAV file - padding (bool): Whether to pad the clip on either side with 1 second of silence - to make sure that short clips can be processed correctly (default: True) + padding (int): How many seconds of silence to pad the start/end of the clip with + to make sure that short clips can be processed correctly (default: 1) kwargs: Any keyword arguments to pass to the class `predict` method Returns: @@ -140,7 +188,13 @@ class Model(): # Load WAV clip frames data = np.frombuffer(f.readframes(f.getnframes()), dtype=np.int16) if padding: - data = np.concatenate((np.zeros(16000).astype(np.int16), data, np.zeros(16000).astype(np.int16))) + data = np.concatenate( + ( + np.zeros(16000*padding).astype(np.int16), + data, + np.zeros(16000*padding).astype(np.int16) + ) + ) # Iterate through clip, getting predictions predictions = [] diff --git a/openwakeword/utils.py b/openwakeword/utils.py index 2b61159..fe3b55a 100644 --- a/openwakeword/utils.py +++ b/openwakeword/utils.py @@ -67,7 +67,7 @@ class AudioFeatures(): self.raw_data_buffer: Deque = deque(maxlen=sr*10) self.melspectrogram_buffer = np.ones((76, 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 = self._get_embeddings(np.zeros(53000).astype(np.int16)) # fill with blank data to start + self.feature_buffer = self._get_embeddings(np.zeros(64000).astype(np.int16)) # fill with blank data to start self.feature_buffer_max_len = 120 # ~10 seconds of feature buffer history def _get_melspectrogram(self, x: Union[np.ndarray, List], melspec_transform: Callable = lambda x: x/10 + 2): diff --git a/pyproject.toml b/pyproject.toml index 5c0733b..8240685 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ requires = ["setuptools>=42"] build-backend = "setuptools.build_meta" [tool.pytest.ini_options] -addopts = "--cov=openwakeword --flake8 --mypy --mypy-ignore-missing-imports" +addopts = "--cov=openwakeword --cov-report term-missing --flake8 --mypy --mypy-ignore-missing-imports" flake8-max-line-length = "120" testpaths = [ "tests", diff --git a/tests/test_models.py b/tests/test_models.py index 53e4bf9..c3bbe74 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -29,11 +29,11 @@ # Imports import openwakeword import os +import numpy as np from pathlib import Path import collections -# Models and corresponding files - +# Define models and corresponding files for testing test_dict = { "hey_mycroft_v1": ["hey_mycroft_v1_test.wav"], "alexa_v5": ["alexa_v5_test.wav"] @@ -43,6 +43,7 @@ test_dict = { # Tests class TestModels: def test_models(self): + # Load model models = [str(i) for i in Path( os.path.join("openwakeword", "resources", "models") ).glob("**/*.onnx") @@ -51,18 +52,60 @@ class TestModels: wakeword_model_paths=models, ) + # 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)) + 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), skipping first prediction as it is innaccurate + # Check scores against default threshold (0.5) for key in predictions_flat.keys(): if key in clip: - assert max(predictions_flat[key][1:]) >= 0.5 + assert max(predictions_flat[key]) >= 0.5 else: - assert max(predictions_flat[key][1:]) < 0.5 + 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") + ).glob("**/*.onnx") + if "embedding" not in str(i) and "melspec" not in str(i)] + owwModel = openwakeword.Model( + wakeword_model_paths=models, + ) + + owwModel.predict(np.zeros(1280), timing=True)