method now supporst numpy array inputs, added more data agmentation options for function

This commit is contained in:
dscripka 2022-12-31 22:19:59 -05:00
parent a3b6985cba
commit f3a8692bda
3 changed files with 50 additions and 23 deletions

View file

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

View file

@ -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 = []

View file

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