Initial version of custom verifier model training and prediction

openwakeword/__init__.py

@@ -1,8 +1,9 @@
 import os
 from openwakeword.model import Model
 from openwakeword.vad import VAD
+from openwakeword.custom_verifier_model import train_custom_verifier
 
-__all__ = ['Model', 'VAD']
+__all__ = ['Model', 'VAD', train_custom_verifier]
 
 models = {
     "alexa": {

openwakeword/custom_verifier_model.py

@@ -0,0 +1,164 @@
+# 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
+from tqdm import tqdm
+import collections
+import openwakeword
+import numpy as np
+import scipy
+import pickle
+
+from sklearn.linear_model import LogisticRegression
+from sklearn.pipeline import make_pipeline
+from sklearn.preprocessing import FunctionTransformer, StandardScaler
+
+# Define functions to prepare data for speaker dependent verifier model
+
+def get_reference_clip_features(
+        reference_clip: str,
+        oww_model: openwakeword.Model,
+        model_name: str,
+        threshold: float = 0.5,
+        N: int = 3,
+        **kwargs
+    ):
+    """
+    Processes input audio files (16-bit, 16-khz single-channel WAV files) and gets the openWakeWord
+    audio features that produce a prediction from the specified model greater than the threshold value.
+    
+
+    Args:
+        reference_clip (str): The target audio file to get features from
+        oww_model (openwakeword.Model): The openWakeWord model object used to get predictions
+        model_name (str): The name of the model to get predictions from (should correspond to
+                          a python dictionary key in the oww_model.models attribute)
+        threshold (float): The minimum score from the model required to capture the associated features
+        N (int): How many times to run feature extraction for a given clip, adding some slight variation
+                 in the starting position each time to ensure that the features are not identical
+
+    Returns:
+        ndarray: A numpy array of shape N x M x L, where N is the number of examples, M is the number
+                 of frames in the window, and L is the audio feature/embedding dimension.
+    """
+
+    # Create dictionary to store frames
+    positive_data = collections.defaultdict(list)
+    
+    # Get predictions
+    for _ in range(N):
+        # Load clip
+        if type(reference_clip) == str:
+            sr, dat = scipy.io.wavfile.read(reference_clip)
+        else:
+            dat = reference_clip
+        
+        # Set random starting point to get small variations in features
+        if N != 1:
+            dat = dat[np.random.randint(0,1280):]
+        
+        # Get predictions
+        step_size = 1280
+        for i in range(0, dat.shape[0]-step_size, step_size):
+            predictions = oww_model.predict(dat[i:i+step_size], **kwargs)
+            if predictions[model_name] >= threshold:
+                features = oww_model.preprocessor.get_features(oww_model.model_inputs[model_name])
+                positive_data[model_name].append(features)
+
+    if len(positive_data[model_name]) == 0:
+        positive_data[model_name].append(np.empty((0, oww_model.model_inputs[model_name], 96)))
+                
+    return np.vstack(positive_data[model_name])
+
+def flatten_features(x):
+    return [i.flatten() for i in x]
+
+def train_verifier_model(features: np.ndarray, labels: np.ndarray):
+    """
+    Train a logistic regression binary classifier model on the provided features and labels
+
+    Args:
+        features (ndarray): A N x M numpy array, where N is the number of examples and M
+                             is the number of features
+        labels (ndarray): A 1D numpy array where each value corresponds to the label of the Nth
+                           example in the `features` argument
+
+    Returns:
+        The trained scikit-learn logistic regression model
+    """
+    # C value matters alot here, depending on dataset size (larger datasets work better with larger C?)
+    clf = LogisticRegression(random_state=0, max_iter=2000, C=0.001)
+    pipeline = make_pipeline(FunctionTransformer(flatten_features), StandardScaler(), clf)
+    pipeline.fit(features, labels)
+
+    return pipeline
+
+def train_custom_verifier(
+        positive_reference_clips: str,
+        negative_reference_clips: str,
+        output_path: str,
+        model_name: str,
+        **kwargs
+        ):
+    """
+    Trains a voice-specific custom verifier model on examples of wake word/phrase speech and other speech
+    from a single user.
+
+    Args:
+        positive_reference_clips (str): The path to a directory containing single-channel 16khz, 16-bit WAV files
+                                        of the target wake word/phrase.
+        negative_reference_clips (str): The path to a directory containing single-channel 16khz, 16-bit WAV files
+                                        of miscellaneous speech not containing the target wake word/phrase.
+        output_path (str): The location to save the trained verifier model (as a scikit-learn .joblib file)
+        model_name (str): The name or path of the trained openWakeWord model that the verifier model will be
+                          based on. If only a name, it must be one of the pre-trained models included in the
+                          openWakeWord release.
+        kwargs: Any other keyword arguments to pass to the openWakeWord model initialization
+
+    Returns:
+        None
+    """
+    # Load target openWakeWord model
+    if os.path.exists(model_name):
+        oww = openwakeword.Model(
+            wakeword_model_paths=[model_name],
+            **kwargs
+        )
+        model_name = model_name[0:-5]
+    else:
+        oww = openwakeword.Model(**kwargs)
+
+    # Get features from positive reference clips
+    positive_features = np.vstack(
+        [get_reference_clip_features(i, oww, model_name, N=5) 
+         for i in tqdm(positive_reference_clips, desc="Processing positive reference clips")]
+    )
+
+    # Get features from negative reference clips
+    negative_features = np.vstack(
+        [get_reference_clip_features(i, oww, model_name, threshold=0.0, N=1) 
+         for i in tqdm(negative_reference_clips, desc="Processing negative reference clips")]
+    )
+    
+    # Train logistic regression model on reference clip features
+    print("Training and saving verifier model...")
+    lr_model = train_verifier_model(
+        np.vstack((positive_features, negative_features)),
+        [1]*positive_features.shape[0] + [0]*negative_features.shape[0]
+    )
+
+    # Save logistic regression model to specified output location
+    print("Done!")
+    pickle.dump(lr_model, open(output_path, "wb"))

openwakeword/model.py

@@ -20,6 +20,7 @@ from openwakeword.utils import AudioFeatures
 
 import wave
 import os
+import pickle
 from collections import deque, defaultdict
 from functools import partial
 import time
@@ -38,6 +39,8 @@ class Model():
             class_mapping_dicts: List[dict] = [],
             enable_speex_noise_suppression: bool = False,
             vad_threshold: float = 0,
+            custom_verifier_models: Union[bool, dict] = False,
+            custom_verifier_threshold: float = 0.1,
             **kwargs
             ):
         """Initialize the openWakeWord model object.
@@ -59,6 +62,13 @@ class Model():
                                    For every input audio frame, a VAD score is obtained and only those model predictions
                                    with VAD scores above the threshold will be returned. The default value (0),
                                    disables voice activity detection entirely.
+            custom_verifier_models (dict): A dictionary of paths to custom verifier models, where 
+                                           the keys are the model names (corresponding to the openwakeword.models attribute)
+                                           and the values are the filepaths of the custom verifier models.
+            custom_verifier_threshold (float): The score threshold to use a custom verifier model. If the score from a model for
+                                               a given frame is greater than this value, the associated custom verifier model will
+                                               also predict on that frame, and the verifier score will be returned.
+            kwargs (dict): Any other keyword arguments to pass the the preprocessor instance
         """
 
         # Initialize the ONNX models and store them
@@ -73,13 +83,16 @@ class Model():
         else:
             wakeword_model_names = [os.path.basename(i[0:-5]) for i in wakeword_model_paths]
 
-        # Create attributes to store models and metadat
+        # Create attributes to store models and metadata
         self.models = {}
         self.model_inputs = {}
         self.model_outputs = {}
         self.class_mapping = {}
         self.model_input_names = {}
+        self.custom_verifier_models = {}
+        self.custom_verifier_threshold = custom_verifier_threshold
         for mdl_path, mdl_name in zip(wakeword_model_paths, wakeword_model_names):
+            # Load openwakeword models
             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]
@@ -92,6 +105,11 @@ class Model():
                 self.class_mapping[mdl_name] = {str(i): str(i) for i in range(0, self.model_outputs[mdl_name])}
             self.model_input_names[mdl_name] = self.models[mdl_name].get_inputs()[0].name
 
+            # Load custom verifier models
+            if isinstance(custom_verifier_models, dict):
+                if custom_verifier_models.get(mdl_name, False):
+                    self.custom_verifier_models[mdl_name] = pickle.load(open(custom_verifier_models[mdl_name], 'rb'))
+
         # Create buffer to store frame predictions
         self.prediction_buffer: DefaultDict[str, deque] = defaultdict(partial(deque, maxlen=30))
 
@@ -183,6 +201,16 @@ class Model():
                 for int_label, cls in self.class_mapping[mdl].items():
                     predictions[cls] = prediction[0][0][int(int_label)]
 
+            # Update scores based on custom verifier model
+            if self.custom_verifier_models != {}:
+                for cls in predictions.keys():
+                    if predictions[cls] >= self.custom_verifier_threshold:
+                        parent_model = self.get_parent_model_from_label(cls)
+                        verifier_prediction = self.custom_verifier_models[parent_model].predict_proba(
+                            self.preprocessor.get_features(self.model_inputs[mdl])
+                        )[0][-1]
+                        predictions[cls] = verifier_prediction
+            
             # Update prediction buffer, and zero predictions for first 5 frames during model initialization
             for cls in predictions.keys():
                 if len(self.prediction_buffer[cls]) < 5: