From f3a8692bdab7d550f46f0c3b869bcc48bfe2e670 Mon Sep 17 00:00:00 2001 From: dscripka Date: Sat, 31 Dec 2022 22:19:59 -0500 Subject: [PATCH] method now supporst numpy array inputs, added more data agmentation options for function --- openwakeword/data.py | 33 +++++++++++++++++++++++---------- openwakeword/model.py | 31 ++++++++++++++++++------------- tests/test_models.py | 9 +++++++++ 3 files changed, 50 insertions(+), 23 deletions(-) diff --git a/openwakeword/data.py b/openwakeword/data.py index ab91293..7c2e764 100644 --- a/openwakeword/data.py +++ b/openwakeword/data.py @@ -260,6 +260,8 @@ def mix_clips_batch( snr_high: float = 0, start_index: List[int] = [], rirs: List[str] = [], + rir_probability: int = 1, + volume_augmentation: bool = True, shuffle: bool = True, seed: int = 0 ): @@ -283,8 +285,13 @@ def mix_clips_batch( 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 + clips to simulate different recording environments. Applies a single random selection from the list RIR file to the entire batch. If empty (the default), nothing is done. + rir_probability (float): The probability (between 0 and 1) that the batch will be convolved with a RIR file. + volume_augmentation (bool): Whether to randomly apply volume augmentation to the clips in the 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). shuffle (bool): Whether to shuffle the foreground clips before mixing (default: True) seed (int): A random seed @@ -342,16 +349,22 @@ def mix_clips_batch( # 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), :] - mixed_clips_batch = reverberate(mixed_clips_batch, rir_waveform, rescale_amp="avg") + if np.random.random() <= rir_probability: + 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), :] + mixed_clips_batch = reverberate(mixed_clips_batch, rir_waveform, rescale_amp="avg") - # Normalize clips only if max value is outside of [-1, 1] - abs_max, _ = torch.max( - torch.abs(mixed_clips_batch), dim=1, keepdim=True - ) - mixed_clips_batch = mixed_clips_batch / abs_max.clamp(min=1.0) + # Apply volume augmentation + if volume_augmentation: + volume_levels = np.random.uniform(0.02, 1.0, mixed_clips_batch.shape[0]) + mixed_clips_batch = (volume_levels/mixed_clips_batch.max(axis=1)[0])[..., None]*mixed_clips_batch + else: + # Normalize clips only if max value is outside of [-1, 1] + abs_max, _ = torch.max( + torch.abs(mixed_clips_batch), dim=1, keepdim=True + ) + mixed_clips_batch = mixed_clips_batch / abs_max.clamp(min=1.0) # Convert to 16-bit PCM audio mixed_clips_batch = (mixed_clips_batch.numpy()*32767).astype(np.int16) diff --git a/openwakeword/model.py b/openwakeword/model.py index a3cc116..9e42c1d 100644 --- a/openwakeword/model.py +++ b/openwakeword/model.py @@ -182,12 +182,13 @@ class Model(): else: return predictions - def predict_clip(self, clip: str, padding: int = 1, **kwargs): + def predict_clip(self, clip: Union[str, np.ndarray], padding: int = 1, **kwargs): """Predict on an full audio clip, simulating streaming prediction. The input clip must bit a 16-bit, 16 khz, single-channel WAV file. Args: - clip (str): The path to a 16-bit PCM, 16 khz, single-channel WAV file + clip (Union[str, np.ndarray]): The path to a 16-bit PCM, 16 khz, single-channel WAV file, + 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) kwargs: Any keyword arguments to pass to the class `predict` method @@ -195,18 +196,22 @@ class Model(): Returns: list: A list containing the frame-level prediction dictionaries for the audio clip """ - # Load audio clip as 16-bit PCM data - with wave.open(clip, mode='rb') as f: - # Load WAV clip frames - data = np.frombuffer(f.readframes(f.getnframes()), dtype=np.int16) - if padding: - data = np.concatenate( - ( - np.zeros(16000*padding).astype(np.int16), - data, - np.zeros(16000*padding).astype(np.int16) - ) + if isinstance(clip, str): + # Load audio clip as 16-bit PCM data + with wave.open(clip, mode='rb') as f: + # Load WAV clip frames + data = np.frombuffer(f.readframes(f.getnframes()), dtype=np.int16) + elif isinstance(clip, np.ndarray): + data = clip + + if padding: + data = np.concatenate( + ( + np.zeros(16000*padding).astype(np.int16), + data, + np.zeros(16000*padding).astype(np.int16) ) + ) # Iterate through clip, getting predictions predictions = [] diff --git a/tests/test_models.py b/tests/test_models.py index e2c41e5..6cc5176 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -71,6 +71,15 @@ class TestModels: else: assert max(predictions_flat[key]) < 0.5 + def test_predict_clip_with_array(self): + # Load model with defaults + owwModel = openwakeword.Model() + + # Make random array and predict + dat = np.random.random(16000) + predictions = owwModel.predict_clip(dat) + assert isinstance(predictions[0], dict) + def test_models_with_timing(self): models = [str(i) for i in Path( os.path.join("openwakeword", "resources", "models")