Implemented patience functionality in prediction, some unit test changes

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)

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 = []

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):

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

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)