All tests (including flake8 and mypy) testing, but coverage is still low

This commit is contained in:
dscripka 2022-10-29 14:44:05 -04:00
parent 2ccb903d61
commit 89230ca697
9 changed files with 295 additions and 190 deletions

View file

@ -1 +1,3 @@
from openwakeword.model import Model
from openwakeword.model import Model
__all__ = ['Model', ]

View file

@ -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)
return np.vstack(X), np.array(y)

View file

@ -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

View file

@ -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
return predictions

View file

@ -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}
return {list(i.keys())[0]: list(i.values())[0] for i in results}

View file

@ -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"
]

Binary file not shown.

Binary file not shown.

View file

@ -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