diff --git a/openwakeword/__init__.py b/openwakeword/__init__.py index cc109d6..904866d 100644 --- a/openwakeword/__init__.py +++ b/openwakeword/__init__.py @@ -1 +1,3 @@ -from openwakeword.model import Model \ No newline at end of file +from openwakeword.model import Model + +__all__ = ['Model', ] diff --git a/openwakeword/data.py b/openwakeword/data.py index f4b5f5d..1918041 100644 --- a/openwakeword/data.py +++ b/openwakeword/data.py @@ -26,8 +26,8 @@ from speechbrain.processing.signal_processing import reverberate import torchaudio import mutagen -# Load audio clips and structure into clips of the same length +# Load audio clips and structure into clips of the same length def stack_clips(audio_data, clip_size=16000*2): """ Takes an input list of 1D arrays (of different lengths), concatenates them together, @@ -51,29 +51,27 @@ def stack_clips(audio_data, clip_size=16000*2): if chunk.shape[0] != clip_size: chunk = np.hstack((chunk, np.zeros(clip_size - chunk.shape[0]))) new_examples.append(chunk) - - # # Convert to 16-bit PCM data - # X = (np.array(new_examples).astype(np.float32)*32767).astype(np.int16) return np.array(new_examples) + def load_audio_clips(files, clip_size=32000): """ Takes the specified audio files and shapes them into an array of N by `clip_size`, where N is determined by the length of the audio files and `clip_size` at run time. - + Clips longer than `clip size` are truncated and extended into the N+1 row. Clips shorter than `clip_size` are combined with the previous or next clip (except for the last clip in `files`, which is ignored if it is too short.) - + Args: files (List[str]): A list of filepaths clip_size (int): The number of samples (of 16khz audio) for all of the rows in the array - + Returns: ndarray: A N by `clip_size` array with the audio data, converted to 16-bit PCM """ - + # Load audio files audio_data = [] for i in files: @@ -85,7 +83,7 @@ def load_audio_clips(files, clip_size=32000): # Get shape of output array N = sum([i.shape[0] for i in audio_data])//clip_size X = np.empty((N, clip_size)) - + # Add audio data to rows previous_row_remainder = None cnt = 0 @@ -95,16 +93,18 @@ def load_audio_clips(files, clip_size=32000): X[cnt, :] = row[0:clip_size] row = row[clip_size:] cnt += 1 - + previous_row_remainder = row if row.size > 0 else None - + # Convert to 16-bit PCM data X = (X*32767).astype(np.int16) return X + # Dato I/O utils + # Convert clips with sox def _convert_clip(input_file, output_file): cmd = f"sox {input_file} -G -r 16000 -c 1 {output_file}" @@ -115,23 +115,23 @@ def _convert_clip(input_file, output_file): def convert_clips(input_files, output_files, sr=16000, ncpu=1): """ Converts files in parallel with multithreading using Sox. - + Intended to only convert input audio files in single-channel, 16 khz clips. - + Args: input_files (List[str]): A list of paths to input files output_files (List[str]): A list of paths ot output files, correspondind 1:1 to the input files sr (int): The output sample rate of the converted clip ncpu (int): The number of CPUs to use for the conversion - + Returns: None """ # Setup ThreadPool object pool = ThreadPool(processes=ncpu) - + # Submit jobs - pool.starmap(_convert_clip, [(i,j) for i,j in zip(input_files, output_files)]) + pool.starmap(_convert_clip, [(i, j) for i, j in zip(input_files, output_files)]) def filter_audio_paths(target_dirs, min_length_secs, max_length_secs, duration_method="size", glob_filter=None): @@ -152,7 +152,7 @@ def filter_audio_paths(target_dirs, min_length_secs, max_length_secs, duration_m much faster, but assumes that all files in the target directory are the same type, sample rate, and bitrate. If None, durations are not calculated. glob_filter (str): A pathlib glob filter string to select specific files within the target directory - + Returns: tuple: A list of strings corresponding to the paths of the wav files that met the length criteria, and a list of their durations (in seconds) @@ -172,7 +172,7 @@ def filter_audio_paths(target_dirs, min_length_secs, max_length_secs, duration_m dir_paths.append(i.path) file_paths.append(i.path) sizes.append(i.stat().st_size) - + if duration_method == "size": durations.extend(estimate_clip_duration(dir_paths, sizes)) @@ -180,16 +180,15 @@ def filter_audio_paths(target_dirs, min_length_secs, max_length_secs, duration_m durations.extend([get_clip_duration(i) for i in tqdm(dir_paths)]) if durations != []: - filtered = [(i,j) for i,j in zip(file_paths, durations) if j >= min_length_secs and j <= max_length_secs] + filtered = [(i, j) for i, j in zip(file_paths, durations) if j >= min_length_secs and j <= max_length_secs] return [i[0] for i in filtered], [i[1] for i in filtered] else: return file_paths, [] - def estimate_clip_duration(audio_files: list, sizes: list): """Estimates the duration of each audio file in a list. - + Assumes that all of the audio files have the same audio format, bit depth, and sample rate. @@ -212,18 +211,20 @@ def estimate_clip_duration(audio_files: list, sizes: list): durations = [] for size in sizes: durations.append((size*8-correction)/details.info.bitrate) - + return durations + def get_clip_duration(clip): """Gets the duration of an audio clip in seconds from file header information""" try: metadata = torchaudio.info(clip) - except RuntimeError: # skip cases where file metadata can't be read + except RuntimeError: # skip cases where file metadata can't be read return 0 return metadata.num_frames/metadata.sample_rate + def get_wav_duration_from_filesize(size, nbytes=2): """ Calculates the duration (in seconds) from a WAV file, assuming it contains 16 khz single-channel audio. @@ -232,7 +233,7 @@ def get_wav_duration_from_filesize(size, nbytes=2): Args: size (int): The file size in bytes nbytes (int): How many bytes for each data point in the audio (e.g., 16-bit is 2, 32-bit is 4, etc.) - + Returns: float: The duration of the audio file in seconds """ @@ -244,30 +245,34 @@ def mix_clips_batch( foreground_clips: List[str], background_clips: List[str], combined_size: int, - batch_size: int=32, - snr_low: float=0, - snr_high: float=0, - start_index: List[int]=[], - rirs: List[str]=[], - shuffle: bool=True, - seed: int=None - ): + batch_size: int = 32, + snr_low: float = 0, + snr_high: float = 0, + start_index: List[int] = [], + rirs: List[str] = [], + shuffle: bool = True, + seed: int = None + ): """ Mixes foreground and background clips at a random SNR level in batches. - + References: https://pytorch.org/audio/main/tutorials/audio_data_augmentation_tutorial.html and https://speechbrain.readthedocs.io/en/latest/API/speechbrain.processing.speech_augmentation.html#speechbrain.processing.speech_augmentation.AddNoise - + Args: foreground_clips (List[str]): A list of paths to the foreground clips - background_clips (List[str]): A list of paths to the background clips (randomly selected for each foreground clip) - combined_size (int): The total length (in samples) of the combined clip. If needed, the background clips are duplicated or truncated to reach this length. + background_clips (List[str]): A list of paths to the background clips (randomly selected for each + foreground clip) + combined_size (int): The total length (in samples) of the combined clip. If needed, the background + clips are duplicated or truncated to reach this length. batch_size (int): The batch size snr_low (float): The low SNR level of the mixing in db snr_high (float): The high snr level of the mixing in db - start_index (List[int]): The starting position (in samples) for the foreground clip to start in the background clip. - rirs (List[str]): A list of paths to room impulse response functions (RIR) to convolve with the clips to simulate different recording environments. - Applies a single random from the list RIR file to the entire batch. If empty (the default), nothing is done. + start_index (List[int]): The starting position (in samples) for the foreground clip to start in + the background clip. + rirs (List[str]): A list of paths to room impulse response functions (RIR) to convolve with the + clips to simulate different recording environments. Applies a single random from the + list RIR file to the entire batch. If empty (the default), nothing is done. shuffle (bool): Whether to shuffle the foreground clips before mixing (default: True) seed (int): A random seed @@ -285,11 +290,11 @@ def mix_clips_batch( # Set start indices, if needed if not start_index: start_index = [0]*batch_size - + for i in range(0, len(foreground_clips), batch_size): # Load foreground clips and truncate (if needed) foreground_clips_batch = [read_audio(i)[0:combined_size] for i in foreground_clips[i:i+batch_size]] - + # Load background clips and pad/truncate as needed background_clips_batch = [read_audio(i) for i in random.sample(background_clips, batch_size)] for ndx, background_clip in enumerate(background_clips_batch): @@ -300,25 +305,25 @@ def mix_clips_batch( elif background_clip.shape[0] > combined_size: r = np.random.randint(0, max(1, background_clip.shape[0] - combined_size)) background_clips_batch[ndx] = background_clip[r:r + combined_size] - + # Mix clips at snr levels snrs_db = np.random.uniform(snr_low, snr_high, batch_size) - mixed_clips_batch = [] + mixed_clips = [] for fg, bg, snr, start in zip(foreground_clips_batch, background_clips_batch, snrs_db, start_index): fg_rms, bg_rms = fg.norm(p=2), bg.norm(p=2) snr = 10 ** (snr / 20) scale = snr * bg_rms / fg_rms bg[start:start + fg.shape[0]] = bg[start:start + fg.shape[0]] + scale*fg[0:bg.shape[0] - start] - mixed_clips_batch.append(bg / 2) - - mixed_clips_batch = torch.vstack(mixed_clips_batch) - + mixed_clips.append(bg / 2) + + mixed_clips_batch = torch.vstack(mixed_clips) + # Apply reverberation to the batch (from a single RIR file) if rirs: rir_waveform, sr = torchaudio.load(random.choice(rirs)) if rir_waveform.shape[0] > 1: - rir_waveform = rir_waveform[random.randint(0,rir_waveform.shape[0]-1), :] + rir_waveform = rir_waveform[random.randint(0, rir_waveform.shape[0]-1), :] mixed_clips_batch = reverberate(mixed_clips_batch, rir_waveform, rescale_amp="avg") # Normalize clips only if max value is outside of [-1, 1] @@ -332,8 +337,8 @@ def mix_clips_batch( yield mixed_clips_batch -# Reverberation data augmentation function +# Reverberation data augmentation function def apply_reverb(x, rir_files): """ Applies reverberation to the input audio clips @@ -354,11 +359,12 @@ def apply_reverb(x, rir_files): # Apply reverberation to the batch (from a single RIR file) if rir_waveform.shape[0] > 1: - rir_waveform = rir_waveform[random.randint(0,rir_waveform.shape[0]-1), :] + rir_waveform = rir_waveform[random.randint(0, rir_waveform.shape[0]-1), :] reverbed = reverberate(torch.from_numpy(x), rir_waveform, rescale_amp="avg") return reverbed.numpy() + # Load batches of data from mmaped numpy arrays class mmap_batch_generator: """ @@ -368,7 +374,13 @@ class mmap_batch_generator: by the `n_per_class` initialization argument. When a mmaped numpy array has been fully interated over, it will restart at the zeroth index automatically. """ - def __init__(self, data_files, batch_size, n_per_class = None, data_transform_funcs = {}, label_transform_funcs = {}): + def __init__(self, + data_files: dict, + batch_size: int, + n_per_class: dict = None, + data_transform_funcs: dict = {}, + label_transform_funcs: dict = {} + ): """ Initialize the generator object @@ -377,36 +389,37 @@ class mmap_batch_generator: Keys should be integer strings representing class labels. batch_size (int): The number of samples per batch n_per_class (dict): A dictionary with integer string labels (as keys) and number of example per batch - (as values). If None (the default), batch sizes for each class will be automatically calculated based on the - the input dataframe shapes and transformation functions. + (as values). If None (the default), batch sizes for each class will be + automatically calculated based on the the input dataframe shapes and transformation + functions. data_transform_funcs (dict): A dictionary of transformation functions to apply to each batch of per class data loaded from the mmaped array. For example, with an array of shape (batch, timesteps, features), if the goal is to half the timesteps per example, (effectively doubling the size of the batch) this function could be passed: - + lambda x: np.vstack( (x[:, 0:timesteps//2, :], x[:, timesteps//2:, :] )) - The user should incorporate the effect of any transform on the values of the + The user should incorporate the effect of any transform on the values of the `n_per_class` argument accordingly, in order to end of with the desired total batch size for each iteration of the generator. label_transform_funcs (dict): A dictionary of transformation functions to apply to each batch of labels. For example, strings can be mapped to integers or one-hot encoded, - groups of classes can be merged together into one, etc. + groups of classes can be merged together into one, etc. """ # inputs self.data_files = data_files self.n_per_class = n_per_class self.data_transform_funcs = data_transform_funcs self.label_transform_funcs = label_transform_funcs - + # Get array mmaps and store their shapes - 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()} - self.shapes = {label:self.data[label].shape for label in self.data.keys()} + 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()} + self.shapes = {label: self.data[label].shape for label in self.data.keys()} # # Update effective shape of mmap array based on user-provided transforms (currently broken) # for lbl, f in self.data_transform_funcs.items(): @@ -431,10 +444,10 @@ class mmap_batch_generator: batches_per_epoch = sum([i[0] for i in self.shapes.values()])//batch_size self.batch_per_epoch = batches_per_epoch print("Batches/steps per epoch:", batches_per_epoch) - + def __iter__(self): return self - + def __next__(self): # Build batch while True: @@ -464,4 +477,4 @@ class mmap_batch_generator: X.append(x) y.extend(y_batch) - return np.vstack(X), np.array(y) \ No newline at end of file + return np.vstack(X), np.array(y) diff --git a/openwakeword/metrics.py b/openwakeword/metrics.py index ee7bb98..268b108 100644 --- a/openwakeword/metrics.py +++ b/openwakeword/metrics.py @@ -12,36 +12,39 @@ # See the License for the specific language governing permissions and # limitations under the License. -## Define metric utility functions specific to the wakeword detection use-case - -from time import time -import matplotlib.pyplot as plt +# Imports import re from tqdm import tqdm import numpy as np -def generate_roc_curve_fprs(scores, n_points=25, time_per_prediction=.08, **kwargs): + +# Define metric utility functions specific to the wakeword detection use-case +def generate_roc_curve_fprs( + scores: list, + n_points: int = 25, + time_per_prediction: float = .08, + **kwargs + ): """ Generates the false positive rate (fpr) per hour for the given predictions over a range of score thresholds. Assumes that all predictions should be less than the threshold, else the prediction is a false positive. Args: - predictions (List): A list of predicted scores, between 0 and 1 - labels (List): A list of ground-truth labels - thresholds (List[float]): A list of threshold values to plot curves for + scores (List): A list of predicted scores, between 0 and 1 + n_points (int): The number of points to use when calculating false positive rates time_per_prediction: The time (in seconds) that each prediction represents Returns: - list: A list of fprs per hour + list: A list of false positive rates per hour at different score threshold levels """ # Determine total time - total_hours = time_per_prediction*len(scores)/3600 # convert to hours + total_hours = time_per_prediction*len(scores)/3600 # convert to hours # Calculate true positive rate fprs = [] - for threshold in tqdm(np.linspace(0.05,0.95,num=n_points)): + for threshold in tqdm(np.linspace(0.05, 0.95, num=n_points)): # Remove repeated predictions from data to not overcount false positives bin_pred = ''.join(["1" if i else "0" for i in np.array(scores) >= threshold]) bin_pred = re.sub("1(0){1,5}1", "1", bin_pred) @@ -50,17 +53,24 @@ def generate_roc_curve_fprs(scores, n_points=25, time_per_prediction=.08, **kwar return fprs -def generate_roc_curve_tprs(scores, n_points=25): + +def generate_roc_curve_tprs( + scores: list, + n_points: int = 25 + ): """ Generates the true positive rate (true accept rate) for the given predictions over a range score thresholds. Assumes that all predictions are supposed to be equal to 1. Args: scores (list): A list of scores for each prediction + + Returns: + list: A list of true positive rates at different score threshold levels """ tprs = [] - for threshold in tqdm(np.linspace(0.05,0.95,num=n_points)): + for threshold in tqdm(np.linspace(0.05, 0.95, num=n_points)): tprs.append(sum(scores >= threshold)/len(scores)) - + return tprs diff --git a/openwakeword/model.py b/openwakeword/model.py index 5e2c0a1..4cacbd8 100644 --- a/openwakeword/model.py +++ b/openwakeword/model.py @@ -21,12 +21,13 @@ import statistics import wave import os from collections import deque, defaultdict -from typing import List from functools import partial import time import pprint -from typing import List, Union +from typing import List, Union, DefaultDict, Dict + +# Define main model class class Model(): """ The main model class for openWakeWord. Creates a model object with the shared audio pre-processer @@ -51,41 +52,44 @@ class Model(): 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.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_input_names[mdl_name] = self.models[mdl_name].get_inputs()[0].name # Create buffer to store frame predictios - self.prediction_buffer = defaultdict(partial(deque, maxlen=5)) + self.prediction_buffer: DefaultDict[str, deque] = defaultdict(partial(deque, maxlen=5)) # Create AudioFeatures object self.preprocessor = AudioFeatures(**kwargs) - - def predict(self, x: Union[np.ndarray, List], median_smooth: bool=False, timing: bool=False): + def predict(self, x: Union[np.ndarray, List], median_smooth: bool = False, 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. - 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. + 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. + 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. Returns: - dict: A dictionary of scores between 0 and 1 for each model, where 0 indicates no wake-word/wake-phrase detected + dict: A dictionary of scores between 0 and 1 for each model, where 0 indicates no + wake-word/wake-phrase detected """ # Get audio features if timing: - timing_dict = {} + timing_dict: Dict[str, Dict] = {} timing_dict["models"] = {} feature_start = time.time() - + self.preprocessor(x) if timing: feature_end = time.time() - timing_dict["preprocessor"] = feature_end - feature_start + timing_dict["models"]["preprocessor"] = feature_end - feature_start # Get predictions from model(s) predictions = {} @@ -113,13 +117,12 @@ class Model(): timing_dict["models"][mdl] = model_end - model_start if timing: - pp = pprint.PrettyPrinter().pprint(timing_dict) + 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: bool = True, **kwargs): """Predict on an full audio clip, simulating streaming prediction. The input clip must bit a 16-bit, 16 khz, single-channel WAV file. @@ -128,7 +131,7 @@ class Model(): 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) kwargs: Any keyword arguments to pass to the class `predict` method - + Returns: list: A list containing the frame-level prediction dictionaries for the audio clip """ @@ -137,7 +140,7 @@ class Model(): # Load WAV clip frames data = np.frombuffer(f.readframes(f.getnframes()), dtype=np.int16) if padding: - data = np.concatenate((np.zeros(32000), data, np.zeros(32000))) + data = np.concatenate((np.zeros(16000).astype(np.int16), data, np.zeros(16000).astype(np.int16))) # Iterate through clip, getting predictions predictions = [] @@ -145,4 +148,4 @@ class Model(): for i in range(0, data.shape[0]-step_size, step_size): predictions.append(self.predict(data[i:i+step_size], **kwargs)) - return predictions \ No newline at end of file + return predictions diff --git a/openwakeword/utils.py b/openwakeword/utils.py index 8db5a9d..2b61159 100644 --- a/openwakeword/utils.py +++ b/openwakeword/utils.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. - # Imports import os import onnxruntime as ort @@ -23,15 +22,37 @@ from multiprocessing.pool import ThreadPool from multiprocessing import Process, Queue import time import openwakeword +from typing import Union, List, Callable, Deque -# Base class for computing audio features using Google's speech_embedding model (https://tfhub.dev/google/speech_embedding/1) + +# Base class for computing audio features using Google's speech_embedding +# model (https://tfhub.dev/google/speech_embedding/1) class AudioFeatures(): + """ + A class for creating audio features from audio data, including melspectograms and Google's + `speech_embedding` features. + """ def __init__(self, - melspec_onnx_model_path=os.path.join(pathlib.Path(__file__).parent.resolve(), "resources", "models", "melspectrogram.onnx"), - embedding_onnx_model_path=os.path.join(pathlib.Path(__file__).parent.resolve(), "resources", "models", "embedding_model.onnx"), - sr=16000, - ncpu=1 - ): + melspec_onnx_model_path: str = os.path.join( + pathlib.Path(__file__).parent.resolve(), + "resources", "models", "melspectrogram.onnx" + ), + embedding_onnx_model_path: str = os.path.join( + pathlib.Path(__file__).parent.resolve(), + "resources", "models", "embedding_model.onnx" + ), + sr: int = 16000, + ncpu: int = 1 + ): + """ + Initialize the AudioFeatures object. + + Args: + melspec_onnx_model_path (str): The path to the ONNX model for computing melspectograms from audio data + embedding_onnx_model_path (str): The path to the ONNX model for Google's `speech_embedding` model + sr (int): The sample rate of the audio (default: 16000 khz) + ncpu (int): The number of CPUs to use when computing melspectrograms and audio features (default: 1) + """ # Initialize the ONNX models sessionOptions = ort.SessionOptions() sessionOptions.inter_op_num_threads = ncpu @@ -43,46 +64,71 @@ class AudioFeatures(): self.onnx_execution_provider = self.melspec_model.get_providers()[0] # Create databuffers - self.raw_data_buffer = 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 feature buffer with blank data to start - self.feature_buffer_max_len = 120 # ~10 seconds of feature buffer history - - def _get_melspectrogram(self, x, melspec_transform = lambda x: x/10 + 2): - """Function to compute the mel-spectrogram of the provided audio samples.""" + 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_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): + """ + Function to compute the mel-spectrogram of the provided audio samples. + + Args: + x (Union[np.ndarray, List]): The input audio data to compute the melspectrogram from + melspec_transform (Callable): A function to transform the computed melspectrogram. Defaults to a transform + that makes the ONNX melspectrogram model closer to the native Tensorflow + implementation from Google. + + Return: + np.ndarray: The computed melspectrogram of the input audio data + """ + # Get input data and adjust type/shape as needed x = np.array(x).astype(np.int16) if isinstance(x, list) else x if x.dtype != np.int16: - raise ValueError(f"Input data must be 16-bit integers (i.e., 16-bit PCM audio). You provided {x.dtype} data.") - x = x[None,] if len(x.shape) < 2 else x - x = x.astype(np.float32) if x.dtype!=np.float32 else x + raise ValueError("Input data must be 16-bit integers (i.e., 16-bit PCM audio)." + f"You provided {x.dtype} data.") + x = x[None, ] if len(x.shape) < 2 else x + x = x.astype(np.float32) if x.dtype != np.float32 else x + + # Get melspectrogram outputs = self.melspec_model.run(None, {'input': x}) spec = np.squeeze(outputs[0]) + if melspec_transform: - spec = melspec_transform(spec) # Arbitrary adjustment to make result closer to original Google speech_embedding model - + spec = melspec_transform(spec) # Arbitrary transform to get result closer to original tensorflow melspec + return spec def _get_embeddings_from_melspec(self, melspec): + """ + Computes the Google `speech_embedding` features from a melspectrogram input + + Args: + melspec (np.ndarray): The input melspectrogram + + Returns: + np.ndarray: The computed audio features/embeddings + """ if melspec.shape[0] != 1: - melspec = melspec[None,] + melspec = melspec[None, ] embedding = self.embedding_model.run(None, {'input_1': melspec})[0].squeeze() return embedding - def _get_embeddings(self, x, window_size=76, step_size = 8, **kwargs): + def _get_embeddings(self, x: np.ndarray, window_size: int = 76, step_size: int = 8, **kwargs): """Function to compute the embeddings of the provide audio samples.""" spec = self._get_melspectrogram(x, **kwargs) 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 + if window.shape[0] == window_size: # truncate short windows windows.append(window) - + batch = np.expand_dims(np.array(windows), axis=-1).astype(np.float32) embedding = self.embedding_model.run(None, {'input_1': batch})[0].squeeze() return embedding - def get_embedding_shape(self, audio_length, sr=16000): + def get_embedding_shape(self, audio_length: float, sr: int = 16000): """Function that determines the size of the output embedding array for a given audio clip length (in seconds)""" x = (np.random.uniform(-1, 1, int(audio_length*sr))*32767).astype(np.int16) return self._get_embeddings(x).shape @@ -90,19 +136,19 @@ class AudioFeatures(): def _get_melspectrogram_batch(self, x, batch_size=128, ncpu=1): """ Compute the melspectrogram of the input audio samples in batches. - - Note that the optimal performance will depend in the interaction between the device, + + Note that the optimal performance will depend in the interaction between the device, batch size, and ncpu (if a CPU device is used). The user is encouraged to experiment with different values of these parameters to identify which combination is best for their data, as often differences of 1-4x are seen. - + Args: x (ndarray): A numpy array of 16 khz input audio data in shape (N, samples). Assumes that all of the audio data is the same length (same number of samples). batch_size (int): The batch size to use when computing the melspectrogram ncpu (int): The number of CPUs to use when computing the melspectrogram. This argument has no effect if the underlying model is executing on a GPU. - + Returns: ndarray: A numpy array of shape (N, frames, melbins) containing the melspectrogram of all N input audio examples @@ -112,49 +158,49 @@ class AudioFeatures(): pool = None if "CPU" in self.onnx_execution_provider: pool = ThreadPool(processes=ncpu) - + # Make batches n_frames = int(np.ceil(x.shape[1]/160-3)) - mel_bins = 32 # fixed by melspectrogram model + mel_bins = 32 # fixed by melspectrogram model melspecs = np.empty((x.shape[0], n_frames, mel_bins), dtype=np.float32) for i in range(0, max(batch_size, x.shape[0]), batch_size): batch = x[i:i+batch_size] - + if "CUDA" in self.onnx_execution_provider: result = self._get_melspectrogram(batch) elif pool: - result = np.array(pool.map(self._get_melspectrogram, batch, chunksize=batch.shape[0]//ncpu)) - + result = np.array(pool.map(self._get_melspectrogram, + batch, chunksize=batch.shape[0]//ncpu)) + melspecs[i:i+batch_size, :, :] = result.squeeze() - + # Cleanup ThreadPool if pool: pool.close() - + return melspecs def _get_embeddings_batch(self, x, batch_size=128, ncpu=1): """ Compute the embeddings of the input melspectrograms in batches. - - Note that the optimal performance will depend in the interaction between the device, + + Note that the optimal performance will depend in the interaction between the device, batch size, and ncpu (if a CPU device is used). The user is encouraged to experiment with different values of these parameters to identify which combination is best for their data, as often differences of 1-4x are seen. - + Args: x (ndarray): A numpy array of melspectrograms of shape (N, frames, melbins). Assumes that all of the melspectrograms have the same shape. batch_size (int): The batch size to use when computing the embeddings ncpu (int): The number of CPUs to use when computing the embeddings. This argument has no effect if the underlying model is executing on a GPU. - + Returns: ndarray: A numpy array of shape (N, frames, embedding_dim) containing the embeddings of all N input melspectrograms """ - # Ensure input is the correct shape if x.shape[1] < 76: raise ValueError("Embedding model requires the input melspectrograms to have at least 76 frames") @@ -163,69 +209,70 @@ class AudioFeatures(): pool = None if "CPU" in self.onnx_execution_provider: pool = ThreadPool(processes=ncpu) - + # Calcuate array sizes and make batches n_frames = (x.shape[1] - 76)//8 + 1 - embedding_dim = 96 # fixed by embedding model + embedding_dim = 96 # fixed by embedding model embeddings = np.empty((x.shape[0], n_frames, embedding_dim), dtype=np.float32) - + batch = [] ndcs = [] for ndx, melspec in enumerate(x): window_size = 76 for i in range(0, melspec.shape[0], 8): window = melspec[i:i+window_size] - if window.shape[0] == window_size: # ignore windows that are too short (truncates end of clip) + if window.shape[0] == window_size: # ignore windows that are too short (truncates end of clip) batch.append(window) ndcs.append(ndx) - + if len(batch) >= batch_size or ndx+1 == x.shape[0]: batch = np.array(batch).astype(np.float32) if "CUDA" in self.onnx_execution_provider: result = self.embedding_model.run(None, {'input_1': batch})[0].squeeze() elif pool: - result = np.array(pool.map(self._get_embeddings_from_melspec, batch, chunksize=batch.shape[0]//ncpu)) - + result = np.array(pool.map(self._get_embeddings_from_melspec, + batch, chunksize=batch.shape[0]//ncpu)) + for j, ndx2 in zip(range(0, result.shape[0], n_frames), ndcs): embeddings[ndx2, :, :] = result[j:j+n_frames] - + batch = [] ndcs = [] - + # Cleanup ThreadPool if pool: pool.close() - + return embeddings def embed_clips(self, x, batch_size=128, ncpu=1): """ Compute the embeddings of the input audio clips in batches. - - Note that the optimal performance will depend in the interaction between the device, + + Note that the optimal performance will depend in the interaction between the device, batch size, and ncpu (if a CPU device is used). The user is encouraged to experiment with different values of these parameters to identify which combination is best for their data, as often differences of 1-4x are seen. - + Args: x (ndarray): A numpy array of 16 khz input audio data in shape (N, samples). Assumes that all of the audio data is the same length (same number of samples). batch_size (int): The batch size to use when computing the embeddings ncpu (int): The number of CPUs to use when computing the melspectrogram. This argument has no effect if the underlying model is executing on a GPU. - + Returns: ndarray: A numpy array of shape (N, frames, embedding_dim) containing the embeddings of all N input audio clips """ - + # Compute melspectrograms melspecs = self._get_melspectrogram_batch(x, batch_size=batch_size, ncpu=ncpu) - + # Compute embeddings from melspectrograms embeddings = self._get_embeddings_batch(melspecs[:, :, :, None], batch_size=batch_size, ncpu=ncpu) - + return embeddings def _streaming_melspectrogram(self, x): @@ -240,37 +287,43 @@ class AudioFeatures(): 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): if len(x) != 1280: - raise ValueError(f"You must provide input samples in frames of 1280 samples @ 1600khz. Received a frame of {len(x)} samples.") + raise ValueError("You must provide input samples in frames of 1280 samples @ 1600khz." + f"Received a frame of {len(x)} samples.") self._streaming_melspectrogram(x) - x = self.melspectrogram_buffer[-76:].astype(np.float32)[None,:,:,None] + x = self.melspectrogram_buffer[-76:].astype(np.float32)[None, :, :, None] if x.shape[1] == 76: - self.feature_buffer = np.vstack((self.feature_buffer, self.embedding_model.run(None, {'input_1': x})[0].squeeze())) - + 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, n_feature_frames=16): - return self.feature_buffer[-n_feature_frames:, :][None,].astype(np.float32) - + return self.feature_buffer[-n_feature_frames:, :][None, ].astype(np.float32) + def __call__(self, x): self._streaming_features(x) # Bulk prediction function -def bulk_predict(file_paths, wakeword_model_paths, input_sizes, ncpu=1, **kwargs): +def bulk_predict( + file_paths: List[str], + wakeword_model_paths: List[str], + ncpu: int = 1, + **kwargs + ): """ Bulk predict on the provided input files in parallel using multiprocessing using the specified model. Args: input_paths (List[str]): The list of input file to predict wakeword_model_path (List[str])): The paths to the wakeword ONNX model files - input_sizes (List[int]): The number of feature columns (e.g., frames) the model expects ncpu (int): How many processes to create (up to max of available CPUs) kwargs (dict): Any other keyword arguments to pass to the model prediction function (`predict_clip`) @@ -288,19 +341,19 @@ def bulk_predict(file_paths, wakeword_model_paths, input_sizes, ncpu=1, **kwargs # Create jobs ps = [] mdls = [] - q = Queue() + q: Queue = Queue() for chunk in chunks: oww = openwakeword.Model( wakeword_model_paths=wakeword_model_paths, - input_sizes=input_sizes, ) mdls.append(oww) + def f(clips): results = [] for clip in clips: results.append({clip: mdls[-1].predict_clip(clip, **kwargs)}) q.put(results) - + ps.append(Process(target=f, args=(chunk,))) # Submit jobs @@ -315,4 +368,4 @@ def bulk_predict(file_paths, wakeword_model_paths, input_sizes, ncpu=1, **kwargs results.extend(q.get()) # Consolidate results and return - return {list(i.keys())[0]:list(i.values())[0] for i in results} \ No newline at end of file + return {list(i.keys())[0]: list(i.values())[0] for i in results} diff --git a/pyproject.toml b/pyproject.toml index b0f0765..5c0733b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,3 +1,11 @@ [build-system] requires = ["setuptools>=42"] build-backend = "setuptools.build_meta" + +[tool.pytest.ini_options] +addopts = "--cov=openwakeword --flake8 --mypy --mypy-ignore-missing-imports" +flake8-max-line-length = "120" +testpaths = [ + "tests", + "openwakeword" +] \ No newline at end of file diff --git a/tests/data/alexa_v5_test.wav b/tests/data/alexa_v5_test.wav new file mode 100644 index 0000000..637b9cd Binary files /dev/null and b/tests/data/alexa_v5_test.wav differ diff --git a/tests/data/hey_mycroft_v1_test.wav b/tests/data/hey_mycroft_v1_test.wav new file mode 100644 index 0000000..8ab6f21 Binary files /dev/null and b/tests/data/hey_mycroft_v1_test.wav differ diff --git a/tests/test_models.py b/tests/test_models.py index afafe54..53e4bf9 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -28,25 +28,41 @@ # Imports import openwakeword -import scipy.io.wavfile -import pytest +import os +from pathlib import Path +import collections + +# Models and corresponding files + +test_dict = { + "hey_mycroft_v1": ["hey_mycroft_v1_test.wav"], + "alexa_v5": ["alexa_v5_test.wav"] +} + # 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] + def test_models(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, ) - step_size = 1280 - 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 + 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)) + + # 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 + for key in predictions_flat.keys(): + if key in clip: + assert max(predictions_flat[key][1:]) >= 0.5 + else: + assert max(predictions_flat[key][1:]) < 0.5