112 lines
3.9 KiB
Python
112 lines
3.9 KiB
Python
# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import random
|
|
import librosa
|
|
from .manifest import Manifest
|
|
from .segment import AudioSegment
|
|
|
|
|
|
class Perturbation(object):
|
|
def max_augmentation_length(self, length):
|
|
return length
|
|
|
|
def perturb(self, data):
|
|
raise NotImplementedError
|
|
|
|
|
|
class SpeedPerturbation(Perturbation):
|
|
def __init__(self, min_speed_rate=0.85, max_speed_rate=1.15, rng=None):
|
|
self._min_rate = min_speed_rate
|
|
self._max_rate = max_speed_rate
|
|
self._rng = random.Random() if rng is None else rng
|
|
|
|
def max_augmentation_length(self, length):
|
|
return length * self._max_rate
|
|
|
|
def perturb(self, data):
|
|
speed_rate = self._rng.uniform(self._min_rate, self._max_rate)
|
|
if speed_rate <= 0:
|
|
raise ValueError("speed_rate should be greater than zero.")
|
|
data._samples = librosa.effects.time_stretch(data._samples, speed_rate)
|
|
|
|
|
|
class GainPerturbation(Perturbation):
|
|
def __init__(self, min_gain_dbfs=-10, max_gain_dbfs=10, rng=None):
|
|
self._min_gain_dbfs = min_gain_dbfs
|
|
self._max_gain_dbfs = max_gain_dbfs
|
|
self._rng = random.Random() if rng is None else rng
|
|
|
|
def perturb(self, data):
|
|
gain = self._rng.uniform(self._min_gain_dbfs, self._max_gain_dbfs)
|
|
data._samples = data._samples * (10. ** (gain / 20.))
|
|
|
|
|
|
|
|
class ShiftPerturbation(Perturbation):
|
|
def __init__(self, min_shift_ms=-5.0, max_shift_ms=5.0, rng=None):
|
|
self._min_shift_ms = min_shift_ms
|
|
self._max_shift_ms = max_shift_ms
|
|
self._rng = random.Random() if rng is None else rng
|
|
|
|
def perturb(self, data):
|
|
shift_ms = self._rng.uniform(self._min_shift_ms, self._max_shift_ms)
|
|
if abs(shift_ms) / 1000 > data.duration:
|
|
# TODO: do something smarter than just ignore this condition
|
|
return
|
|
shift_samples = int(shift_ms * data.sample_rate // 1000)
|
|
# print("DEBUG: shift:", shift_samples)
|
|
if shift_samples < 0:
|
|
data._samples[-shift_samples:] = data._samples[:shift_samples]
|
|
data._samples[:-shift_samples] = 0
|
|
elif shift_samples > 0:
|
|
data._samples[:-shift_samples] = data._samples[shift_samples:]
|
|
data._samples[-shift_samples:] = 0
|
|
|
|
|
|
perturbation_types = {
|
|
"speed": SpeedPerturbation,
|
|
"gain": GainPerturbation,
|
|
"shift": ShiftPerturbation,
|
|
}
|
|
|
|
|
|
class AudioAugmentor(object):
|
|
def __init__(self, perturbations=None, rng=None):
|
|
self._rng = random.Random() if rng is None else rng
|
|
self._pipeline = perturbations if perturbations is not None else []
|
|
|
|
def perturb(self, segment):
|
|
for (prob, p) in self._pipeline:
|
|
if self._rng.random() < prob:
|
|
p.perturb(segment)
|
|
return
|
|
|
|
def max_augmentation_length(self, length):
|
|
newlen = length
|
|
for (prob, p) in self._pipeline:
|
|
newlen = p.max_augmentation_length(newlen)
|
|
return newlen
|
|
|
|
@classmethod
|
|
def from_config(cls, config):
|
|
ptbs = []
|
|
for p in config:
|
|
if p['aug_type'] not in perturbation_types:
|
|
print(p['aug_type'], "perturbation not known. Skipping.")
|
|
continue
|
|
perturbation = perturbation_types[p['aug_type']]
|
|
ptbs.append((p['prob'], perturbation(**p['cfg'])))
|
|
return cls(perturbations=ptbs)
|