akadmin/openwakeword
1
0
Fork 0
mirror of https://github.com/dscripka/openWakeWord.git synced 2026-08-27 18:17:20 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

Integrated Speex noise suppression as an optional pre-processing step, adjust data augmentation method.

Browse source
This commit is contained in:
David Scripka 2023-01-14 19:47:54 -05:00
parent 81badf5728
commit f870d4340d
3 changed files with 89 additions and 42 deletions
Download patch file Download diff file Expand all files Collapse all files

6
openwakeword/__init__.py
View file

@ -1,7 +1,8 @@
import os
from openwakeword.model import Model
from openwakeword.vad import VAD
__all__ = ['Model', ]
__all__ = ['Model', 'VAD']
models = {
"alexa": {
@ -10,6 +11,9 @@ models = {
"hey_mycroft": {
"model_path": os.path.join(os.path.dirname(os.path.abspath(__file__)), "resources/models/hey_mycroft_v0.1.onnx")
},
"hey_jarvis": {
"model_path": os.path.join(os.path.dirname(os.path.abspath(__file__)), "resources/models/hey_jarvis_v0.1.onnx")
},
"timer": {
"model_path": os.path.join(os.path.dirname(os.path.abspath(__file__)), "resources/models/timer_v0.1.onnx")
},

81
openwakeword/data.py
View file

@ -304,9 +304,9 @@ def mix_clips_batch(
This simply scales the data of each clip such that the maximum value is is between
0.02 and 1.0 (the floor shouldn't be zero as beyond a certain point the audio data
is no longer valid).
generated_noise_augmentation: The probability of mixing the foreground clip with generated random noise
instead of a provided background clip. With be either "white", "brown",
"blue", "pink", or "violet" noise.
generated_noise_augmentation: The probability of further mixing the mixed clip with generated random noise.
Will be either "white", "brown", "blue", "pink", or "violet" noise, mixed at a
random SNR between `snr_low` and `snr_high`.
return_background_clips (bool): Whether to return the segment of the background clip that was mixed with each
foreground clip in the batch.
return_background_clips_delay (Tuple(int)): The lower and upper bound of a random delay (in samples)
@ -359,42 +359,38 @@ def mix_clips_batch(
labels_batch = np.array(labels[i:i+batch_size])
# Load background clips and pad/truncate as needed
if np.random.random() < generated_noise_augmentation:
noise_color = ["white", "pink", "blue", "brown", "violet"]
background_clips_batch = [torch.from_numpy(
acoustics.generator.noise(combined_size, color=np.random.choice(noise_color))
)
for _ in range(batch_size)]
background_clips_batch = [i/i.max() for i in background_clips_batch]
else:
background_clips_batch = [read_audio(j) for j in random.sample(background_clips, batch_size)]
background_clips_batch = [j[0] if len(j.shape) > 1 else j for j in background_clips_batch]
background_clips_batch_delayed = []
delay = np.random.randint(return_background_clips_delay[0], return_background_clips_delay[1] + 1)
for ndx, background_clip in enumerate(background_clips_batch):
if background_clip.shape[0] < (combined_size + delay):
repeated = background_clip.repeat(
np.ceil((combined_size + delay)/background_clip.shape[0]).astype(np.int32)
)
background_clips_batch[ndx] = repeated[0:combined_size]
background_clips_batch_delayed.append(repeated[0+delay:combined_size + delay].clone())
elif background_clip.shape[0] > (combined_size + delay):
r = np.random.randint(0, max(1, background_clip.shape[0] - combined_size - delay))
background_clips_batch[ndx] = background_clip[r:r + combined_size]
background_clips_batch_delayed.append(background_clip[r+delay:r + combined_size + delay].clone())
background_clips_batch = [read_audio(j) for j in random.sample(background_clips, batch_size)]
background_clips_batch = [j[0] if len(j.shape) > 1 else j for j in background_clips_batch]
background_clips_batch_delayed = []
delay = np.random.randint(return_background_clips_delay[0], return_background_clips_delay[1] + 1)
for ndx, background_clip in enumerate(background_clips_batch):
if background_clip.shape[0] < (combined_size + delay):
repeated = background_clip.repeat(
np.ceil((combined_size + delay)/background_clip.shape[0]).astype(np.int32)
)
background_clips_batch[ndx] = repeated[0:combined_size]
background_clips_batch_delayed.append(repeated[0+delay:combined_size + delay].clone())
elif background_clip.shape[0] > (combined_size + delay):
r = np.random.randint(0, max(1, background_clip.shape[0] - combined_size - delay))
background_clips_batch[ndx] = background_clip[r:r + combined_size]
background_clips_batch_delayed.append(background_clip[r+delay:r + combined_size + delay].clone())
# Mix clips at snr levels
snrs_db = np.random.uniform(snr_low, snr_high, batch_size)
mixed_clips = []
for fg, bg, snr, start in zip(foreground_clips_batch, background_clips_batch,
snrs_db, start_index_batch):
fg_rms, bg_rms = fg.norm(p=2), bg.norm(p=2)
snr = 10 ** (snr / 20)
scale = snr * bg_rms / fg_rms
if bg.shape[0] != combined_size:
raise ValueError(bg.shape)
bg[start:start + fg.shape[0]] = bg[start:start + fg.shape[0]] + scale*fg
mixed_clips.append(bg / 2)
mixed_clip = mix_clip(fg, bg, snr, start)
if np.random.random() < generated_noise_augmentation:
noise_color = ["white", "pink", "blue", "brown", "violet"]
noise_clip = acoustics.generator.noise(combined_size, color=np.random.choice(noise_color))
noise_clip = torch.from_numpy(noise_clip/noise_clip.max())
mixed_clip = mix_clip(mixed_clip, noise_clip, np.random.choice(snrs_db), 0)
mixed_clips.append(mixed_clip)
mixed_clips_batch = torch.vstack(mixed_clips)
@ -432,6 +428,13 @@ def mix_clips_batch(
* 32767).astype(np.int16)[error_index]
yield mixed_clips_batch, labels_batch, background_clips_batch_delayed
def mix_clip(fg, bg, snr, start):
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
return bg / 2
def truncate_clip(x, max_size, method="truncate_start"):
"""
Truncates and audio clip with the specified method
@ -500,7 +503,7 @@ class mmap_batch_generator:
"""
def __init__(self,
data_files: dict,
batch_size: int,
batch_size: int = 128,
n_per_class: dict = {},
data_transform_funcs: dict = {},
label_transform_funcs: dict = {}
@ -552,22 +555,22 @@ class mmap_batch_generator:
# self.shapes[lbl] = (new_shape[0]*self.original_shapes[lbl][0], new_shape[1], new_shape[2])
# Calculate batch sizes, if the user didn't specify them
scale_factor = 1
if not self.n_per_class:
self.n_per_class = {}
for lbl, shape in self.shapes.items():
dummy_data = np.random.random((10, self.shapes[lbl][1], self.shapes[lbl][2]))
if self.data_transform_funcs.get(lbl, None):
scale_factor = self.data_transform_funcs.get(lbl, None)(dummy_data).shape[0]/10
else:
scale_factor = 1
ratio = self.shapes[lbl][0]/sum([i[0] for i in self.shapes.values()])
self.n_per_class[lbl] = max(1, int(int(batch_size*ratio)/scale_factor))
# Get estimated batches per epoch, including the effect of any user-provided transforms
batch_size = sum([val*scale_factor for val in self.n_per_class.values()])
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)
# Get estimated batches per epoch, including the effect of any user-provided transforms
batch_size = sum([val*scale_factor for val in self.n_per_class.values()])
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

44
openwakeword/model.py
View file

@ -36,6 +36,7 @@ class Model():
self,
wakeword_model_paths: List[str] = [],
class_mapping_dicts: List[dict] = [],
enable_speex_noise_suppression: bool = False,
**kwargs
):
"""Initialize the openWakeWord model object.
@ -46,6 +47,10 @@ class Model():
class_mapping_dicts (List[dict]): A list of dictionaries with integer to string class mappings for
each model in the `wakeword_model_paths` arguments
(e.g., {"0": "class_1", "1": "class_2"})
enable_speex_noise_suppression (bool): Whether to use the noise suppresion from the SpeexDSP library to pre-process
all incoming audio. May increase model performance when reasonably stationary
background noise is present in the environment where openWakeWord will be used.
It is very lightweight, so enabling it doesn't significantly impact efficiency.
"""
# Initialize the ONNX models and store them
@ -82,6 +87,13 @@ class Model():
# Create buffer to store frame predictions
self.prediction_buffer: DefaultDict[str, deque] = defaultdict(partial(deque, maxlen=30))
# Initialize SpeexDSP noise canceller
if enable_speex_noise_suppression:
from speexdsp_ns import NoiseSuppression
self.speex_ns = NoiseSuppression.create(160, 16000)
else:
self.speex_ns = None
# Create AudioFeatures object
self.preprocessor = AudioFeatures(**kwargs)
@ -123,13 +135,16 @@ class Model():
wake-word/wake-phrase detected. If the `timing` argument is true, returns a
tuple of dicts containing model predictions and timing information, respectively.
"""
# Get audio features
# Get audio features (optionally with Speex noise suppression)
if timing:
timing_dict: Dict[str, Dict] = {}
timing_dict["models"] = {}
feature_start = time.time()
self.preprocessor(x)
if self.speex_ns:
self.preprocessor(self._suppress_noise_with_speex(x))
else:
self.preprocessor(x)
if timing:
feature_end = time.time()
@ -220,3 +235,28 @@ class Model():
predictions.append(self.predict(data[i:i+step_size], **kwargs))
return predictions
def _suppress_noise_with_speex(self, x: np.ndarray, frame_size: int = 160):
"""
Runs the input audio through the SpeexDSP noise suppression algorithm.
Note that this function updates the state of the existing Speex noise
suppression object, and isn't intended to be called externally.
Args:
x (ndarray): The 16-bit, 16khz audio to process. Must always be an
integer multiple of `frame_size`.
frame_size (int): The frame size to use for the Speex Noise suppressor.
Must match the frame size specified during the
initialization of the noise suppressor.
Returns:
ndarray: The input audio with noise suppression applied
"""
cleaned = []
for i in range(0, x.shape[0], frame_size):
chunk = x[i:i+frame_size]
cleaned.append(self.speex_ns.process(chunk.tobytes()))
cleaned = b''.join(cleaned)
cleaned = np.frombuffer(cleaned, np.int16)
return cleaned