Changed overall prediction behavior to support varying sizes of audio frames/chunks, which leads to efficiency gains on streaming melspectrogram calculations

This commit is contained in:
dscripka 2023-03-20 21:32:39 -04:00
parent 8322a96fa8
commit cf2ebc176d
4 changed files with 88 additions and 29 deletions

View file

@ -16,27 +16,40 @@
import pyaudio
import numpy as np
from openwakeword.model import Model
import argparse
# Parse input arguments
parser=argparse.ArgumentParser()
parser.add_argument(
"--chunk_size",
help="How much audio (in samples) to predict on at once",
type=int,
default=1280,
required=True
)
args=parser.parse_args()
# Get microphone stream
FORMAT = pyaudio.paInt16
CHANNELS = 1
RATE = 16000
CHUNK = 1280
CHUNK = args.chunk_size
audio = pyaudio.PyAudio()
mic_stream = audio.open(format=FORMAT, channels=CHANNELS, rate=RATE, input=True, frames_per_buffer=CHUNK)
# Load pre-trained openwakeword models
owwModel = Model()
# Run capture loop, checking for hotwords
# Run capture loop continuosly, checking for wakewords
if __name__ == "__main__":
# Predict continuously on audio stream
# Generate output string header
print("\n\n")
print("#"*100)
print("Listening for wakewords...")
print("#"*100)
print("\n"*13)
while True:
# Get audio
audio = np.frombuffer(mic_stream.read(CHUNK), dtype=np.int16)
@ -44,7 +57,7 @@ if __name__ == "__main__":
# Feed to openWakeWord model
prediction = owwModel.predict(audio)
# Generate output string header
# Column titles
n_spaces = 16
output_string_header = """
Model Name | Score | Wakeword Status

View file

@ -200,11 +200,26 @@ class Model():
if timing:
model_start = time.time()
# Run model
prediction = self.models[mdl].run(
None,
{input_name: self.preprocessor.get_features(self.model_inputs[mdl])}
)
# Run model to get predictions
if len(x) > 1280:
group_predictions = []
for i in np.arange(len(x)//1280-1, -1, -1):
group_predictions.extend(
self.models[mdl].run(
None,
{input_name: self.preprocessor.get_features(
self.model_inputs[mdl],
start_ndx=-self.model_inputs[mdl] - i
)}
)
)
prediction = np.array(group_predictions).max(axis=0)[None, ]
else:
prediction = self.models[mdl].run(
None,
{input_name: self.preprocessor.get_features(self.model_inputs[mdl])}
)
if self.model_outputs[mdl] == 1:
predictions[mdl] = prediction[0][0][0]
else:
@ -267,7 +282,7 @@ class Model():
else:
return predictions
def predict_clip(self, clip: Union[str, np.ndarray], padding: int = 1, **kwargs):
def predict_clip(self, clip: Union[str, np.ndarray], padding: int = 1, chunk_size=1280, **kwargs):
"""Predict on an full audio clip, simulating streaming prediction.
The input clip must bit a 16-bit, 16 khz, single-channel WAV file.
@ -276,6 +291,7 @@ class Model():
or an 1D array containing the same type of data
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)
chunk_size (int): The size (in samples) of each chunk of audio to pass to the model
kwargs: Any keyword arguments to pass to the class `predict` method
Returns:
@ -300,7 +316,7 @@ class Model():
# Iterate through clip, getting predictions
predictions = []
step_size = 1280
step_size = chunk_size
for i in range(0, data.shape[0]-step_size, step_size):
predictions.append(self.predict(data[i:i+step_size], **kwargs))

View file

@ -67,6 +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.accumulated_samples = 0 # the samples added to the buffer since the audio preprocessor was last called
self.feature_buffer = self._get_embeddings(np.zeros(160000).astype(np.int16)) # fill with blank data to start
self.feature_buffer_max_len = 120 # ~10 seconds of feature buffer history
@ -210,7 +211,7 @@ class AudioFeatures():
if "CPU" in self.onnx_execution_provider:
pool = ThreadPool(processes=ncpu)
# Calcuate array sizes and make batches
# Calculate array sizes and make batches
n_frames = (x.shape[1] - 76)//8 + 1
embedding_dim = 96 # fixed by embedding model
embeddings = np.empty((x.shape[0], n_frames, embedding_dim), dtype=np.float32)
@ -275,37 +276,62 @@ class AudioFeatures():
return embeddings
def _streaming_melspectrogram(self, x):
def _streaming_melspectrogram(self, n_samples):
"""Note! There seem to be some slight numerical issues depending on the underlying audio data
such that the streaming method is not exactly the same as when the melspectrogram of the entire
clip is calculated. It's unclear if this difference is significant and will impact model performance.
In particular padding with 0 or very small values seems to demonstrate the differences well.
"""
if len(x) < 400:
raise ValueError("The number of input frames must be at least 400 samples @ 16khz (25 ms)!")
self.raw_data_buffer.extend(x.tolist() if isinstance(x, np.ndarray) else x)
self.melspectrogram_buffer = np.vstack(
(self.melspectrogram_buffer, self._get_melspectrogram(list(self.raw_data_buffer)[-len(x)-160*3:]))
(self.melspectrogram_buffer, self._get_melspectrogram(list(self.raw_data_buffer)[-n_samples-160*3:]))
)
if self.melspectrogram_buffer.shape[0] > self.melspectrogram_max_len:
self.melspectrogram_buffer = self.melspectrogram_buffer[-self.melspectrogram_max_len:, :]
def _buffer_raw_data(self, x):
"""
Adds raw audio data to the input buffer
"""
if len(x) < 400:
raise ValueError("The number of input frames must be at least 400 samples @ 16khz (25 ms)!")
self.raw_data_buffer.extend(x.tolist() if isinstance(x, np.ndarray) else x)
def _streaming_features(self, x):
if len(x) != 1280:
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]
if x.shape[1] == 76:
self.feature_buffer = np.vstack((self.feature_buffer,
self.embedding_model.run(None, {'input_1': x})[0].squeeze()))
# if len(x) != 1280:
# raise ValueError("You must provide input samples in frames of 1280 samples @ 1600khz."
# f"Received a frame of {len(x)} samples.")
# Add raw audio data to buffer
self._buffer_raw_data(x)
self.accumulated_samples += len(x)
# Only calculate melspectrogram every ~0.5 seconds to significantly increase efficiency
if self.accumulated_samples >= 1280:
self._streaming_melspectrogram(self.accumulated_samples)
# Calculate new audio embeddings/features based on update melspectrograms
for i in np.arange(self.accumulated_samples//1280-1, -1, -1):
ndx = -8*i
ndx = ndx if ndx != 0 else len(self.melspectrogram_buffer)
x = self.melspectrogram_buffer[-76 + ndx:ndx].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()))
# Reset raw data buffer counter
self.accumulated_samples = 0
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: int = 16):
return self.feature_buffer[int(-1*n_feature_frames):, :][None, ].astype(np.float32)
def get_features(self, n_feature_frames: int = 16, start_ndx: int = -1):
if start_ndx != -1:
end_ndx = start_ndx + int(n_feature_frames) \
if start_ndx + n_feature_frames != 0 else len(self.feature_buffer)
return self.feature_buffer[start_ndx:end_ndx, :][None, ].astype(np.float32)
else:
return self.feature_buffer[int(-1*n_feature_frames):, :][None, ].astype(np.float32)
def __call__(self, x):
self._streaming_features(x)
@ -328,7 +354,8 @@ def bulk_predict(
prediction_function (str): The name of the method used to predict on the input audio files
(default is the `predict_clip` method)
ncpu (int): How many processes to create (up to max of available CPUs)
kwargs (dict): Any other keyword arguments to pass to the model initialization
kwargs (dict): Any other keyword arguments to pass to the model initialization or
specified prediction function
Returns:
dict: A dictionary containing the predictions for each file, with the filepath as the key

View file

@ -47,6 +47,9 @@ class TestModels:
# Prediction on random data
owwModel.predict(np.random.randint(-1000, 1000, 1280).astype(np.int16))
# Prediction on random data with different chunk size
owwModel.predict(np.random.randint(-1000, 1000, 1280*2).astype(np.int16))
def test_custom_model_label_mapping_dict(self):
# Load model with model path
owwModel = openwakeword.Model(wakeword_model_paths=[