import typing
import julius
import numpy as np
import torch
import torchaudio
from . import util
class EffectMixin:
GAIN_FACTOR = np.log(10) / 20
"""Gain factor for converting between amplitude and decibels."""
CODEC_PRESETS = {
"8-bit": {"format": "wav", "encoding": "ULAW", "bits_per_sample": 8},
"GSM-FR": {"format": "gsm"},
"MP3": {"format": "mp3", "compression": -9},
"Vorbis": {"format": "vorbis", "compression": -1},
"Ogg": {
"format": "ogg",
"compression": -1,
},
"Amr-nb": {"format": "amr-nb"},
}
"""Presets for applying codecs via torchaudio."""
def mix(
self,
other,
snr: typing.Union[torch.Tensor, np.ndarray, float] = 10,
other_eq: typing.Union[torch.Tensor, np.ndarray] = None,
):
"""Mixes noise with signal at specified
signal-to-noise ratio. Optionally, the
other signal can be equalized in-place.
Parameters
----------
other : AudioSignal
AudioSignal object to mix with.
snr : typing.Union[torch.Tensor, np.ndarray, float], optional
Signal to noise ratio, by default 10
other_eq : typing.Union[torch.Tensor, np.ndarray], optional
EQ curve to apply to other signal, if any, by default None
Returns
-------
AudioSignal
In-place modification of AudioSignal.
"""
snr = util.ensure_tensor(snr).to(self.device)
pad_len = max(0, self.signal_length - other.signal_length)
other.zero_pad(0, pad_len)
other.truncate_samples(self.signal_length)
if other_eq is not None:
other = other.equalizer(other_eq)
tgt_loudness = self.loudness() - snr
other = other.normalize(tgt_loudness)
self.audio_data = self.audio_data + other.audio_data
return self
def convolve(self, other, start_at_max: bool = True):
"""Convolves self with other.
This function uses FFTs to do the convolution.
Parameters
----------
other : AudioSignal
Signal to convolve with.
start_at_max : bool, optional
Whether to start at the max value of other signal, to
avoid inducing delays, by default True
Returns
-------
AudioSignal
Convolved signal, in-place.
"""
from . import AudioSignal
pad_len = self.signal_length - other.signal_length
if pad_len > 0:
other.zero_pad(0, pad_len)
else:
other.truncate_samples(self.signal_length)
if start_at_max:
# Use roll to rotate over the max for every item
# so that the impulse responses don't induce any
# delay.
idx = other.audio_data.abs().argmax(axis=-1)
irs = torch.zeros_like(other.audio_data)
for i in range(other.batch_size):
irs[i] = torch.roll(other.audio_data[i], -idx[i].item(), -1)
other = AudioSignal(irs, other.sample_rate)
delta = torch.zeros_like(other.audio_data)
delta[..., 0] = 1
length = self.signal_length
delta_fft = torch.fft.rfft(delta, length)
other_fft = torch.fft.rfft(other.audio_data, length)
self_fft = torch.fft.rfft(self.audio_data, length)
convolved_fft = other_fft * self_fft
convolved_audio = torch.fft.irfft(convolved_fft, length)
delta_convolved_fft = other_fft * delta_fft
delta_audio = torch.fft.irfft(delta_convolved_fft, length)
# Use the delta to rescale the audio exactly as needed.
delta_max = delta_audio.abs().max(dim=-1, keepdims=True)[0]
scale = 1 / delta_max.clamp(1e-5)
convolved_audio = convolved_audio * scale
self.audio_data = convolved_audio
return self
def apply_ir(
self,
ir,
drr: typing.Union[torch.Tensor, np.ndarray, float] = None,
ir_eq: typing.Union[torch.Tensor, np.ndarray] = None,
use_original_phase: bool = False,
):
"""Applies an impulse response to the signal. If ` is`ir_eq``
is specified, the impulse response is equalized before
it is applied, using the given curve.
Parameters
----------
ir : AudioSignal
Impulse response to convolve with.
drr : typing.Union[torch.Tensor, np.ndarray, float], optional
Direct-to-reverberant ratio that impulse response will be
altered to, if specified, by default None
ir_eq : typing.Union[torch.Tensor, np.ndarray], optional
Equalization that will be applied to impulse response
if specified, by default None
use_original_phase : bool, optional
Whether to use the original phase, instead of the convolved
phase, by default False
Returns
-------
AudioSignal
Signal with impulse response applied to it
"""
if ir_eq is not None:
ir = ir.equalizer(ir_eq)
if drr is not None:
ir = ir.alter_drr(drr)
# Save the peak before
max_spk = self.audio_data.abs().max(dim=-1, keepdims=True).values
# Augment the impulse response to simulate microphone effects
# and with varying direct-to-reverberant ratio.
phase = self.phase
self.convolve(ir)
# Use the input phase
if use_original_phase:
self.stft()
self.stft_data = self.magnitude * torch.exp(1j * phase)
self.istft()
# Rescale to the input's amplitude
max_transformed = self.audio_data.abs().max(dim=-1, keepdims=True).values
scale_factor = max_spk.clamp(1e-8) / max_transformed.clamp(1e-8)
self = self * scale_factor
return self
def ensure_max_of_audio(self, max: float = 1.0):
"""Ensures that ``abs(audio_data) <= max``.
Parameters
----------
max : float, optional
Max absolute value of signal, by default 1.0
Returns
-------
AudioSignal
Signal with values scaled between -max and max.
"""
peak = self.audio_data.abs().max(dim=-1, keepdims=True)[0]
peak_gain = torch.ones_like(peak)
peak_gain[peak > max] = max / peak[peak > max]
self.audio_data = self.audio_data * peak_gain
return self
def normalize(self, db: typing.Union[torch.Tensor, np.ndarray, float] = -24.0):
"""Normalizes the signal's volume to the specified db, in LUFS.
This is GPU-compatible, making for very fast loudness normalization.
Parameters
----------
db : typing.Union[torch.Tensor, np.ndarray, float], optional
Loudness to normalize to, by default -24.0
Returns
-------
AudioSignal
Normalized audio signal.
"""
db = util.ensure_tensor(db).to(self.device)
ref_db = self.loudness()
gain = db - ref_db
gain = torch.exp(gain * self.GAIN_FACTOR)
self.audio_data = self.audio_data * gain[:, None, None]
return self
def volume_change(self, db: typing.Union[torch.Tensor, np.ndarray, float]):
"""Change volume of signal by some amount, in dB.
Parameters
----------
db : typing.Union[torch.Tensor, np.ndarray, float]
Amount to change volume by.
Returns
-------
AudioSignal
Signal at new volume.
"""
db = util.ensure_tensor(db, ndim=1).to(self.device)
gain = torch.exp(db * self.GAIN_FACTOR)
self.audio_data = self.audio_data * gain[:, None, None]
return self
def _to_2d(self):
waveform = self.audio_data.reshape(-1, self.signal_length)
return waveform
def _to_3d(self, waveform):
return waveform.reshape(self.batch_size, self.num_channels, -1)
def pitch_shift(self, n_semitones: int, quick: bool = True):
"""Pitch shift the signal. All items in the batch
get the same pitch shift.
Parameters
----------
n_semitones : int
How many semitones to shift the signal by.
quick : bool, optional
Using quick pitch shifting, by default True
Returns
-------
AudioSignal
Pitch shifted audio signal.
"""
device = self.device
effects = [
["pitch", str(n_semitones * 100)],
["rate", str(self.sample_rate)],
]
if quick:
effects[0].insert(1, "-q")
waveform = self._to_2d().cpu()
waveform, sample_rate = torchaudio.sox_effects.apply_effects_tensor(
waveform, self.sample_rate, effects, channels_first=True
)
self.sample_rate = sample_rate
self.audio_data = self._to_3d(waveform)
return self.to(device)
def time_stretch(self, factor: float, quick: bool = True):
"""Time stretch the audio signal.
Parameters
----------
factor : float
Factor by which to stretch the AudioSignal. Typically
between 0.8 and 1.2.
quick : bool, optional
Whether to use quick time stretching, by default True
Returns
-------
AudioSignal
Time-stretched AudioSignal.
"""
device = self.device
effects = [
["tempo", str(factor)],
["rate", str(self.sample_rate)],
]
if quick:
effects[0].insert(1, "-q")
waveform = self._to_2d().cpu()
waveform, sample_rate = torchaudio.sox_effects.apply_effects_tensor(
waveform, self.sample_rate, effects, channels_first=True
)
self.sample_rate = sample_rate
self.audio_data = self._to_3d(waveform)
return self.to(device)
def apply_codec(
self,
preset: str = None,
format: str = "wav",
encoding: str = None,
bits_per_sample: int = None,
compression: int = None,
): # pragma: no cover
"""Applies an audio codec to the signal.
Parameters
----------
preset : str, optional
One of the keys in ``self.CODEC_PRESETS``, by default None
format : str, optional
Format for audio codec, by default "wav"
encoding : str, optional
Encoding to use, by default None
bits_per_sample : int, optional
How many bits per sample, by default None
compression : int, optional
Compression amount of codec, by default None
Returns
-------
AudioSignal
AudioSignal with codec applied.
Raises
------
ValueError
If preset is not in ``self.CODEC_PRESETS``, an error
is thrown.
"""
torchaudio_version_070 = "0.7" in torchaudio.__version__
if torchaudio_version_070:
return self
kwargs = {
"format": format,
"encoding": encoding,
"bits_per_sample": bits_per_sample,
"compression": compression,
}
if preset is not None:
if preset in self.CODEC_PRESETS:
kwargs = self.CODEC_PRESETS[preset]
else:
raise ValueError(
f"Unknown preset: {preset}. "
f"Known presets: {list(self.CODEC_PRESETS.keys())}"
)
waveform = self._to_2d()
if kwargs["format"] in ["vorbis", "mp3", "ogg", "amr-nb"]:
# Apply it in a for loop
augmented = torch.cat(
[
torchaudio.functional.apply_codec(
waveform[i][None, :], self.sample_rate, **kwargs
)
for i in range(waveform.shape[0])
],
dim=0,
)
else:
augmented = torchaudio.functional.apply_codec(
waveform, self.sample_rate, **kwargs
)
augmented = self._to_3d(augmented)
self.audio_data = augmented
return self
def mel_filterbank(self, n_bands: int):
"""Breaks signal into mel bands.
Parameters
----------
n_bands : int
Number of mel bands to use.
Returns
-------
torch.Tensor
Mel-filtered bands, with last axis being the band index.
"""
filterbank = (
julius.SplitBands(self.sample_rate, n_bands).float().to(self.device)
)
filtered = filterbank(self.audio_data)
return filtered.permute(1, 2, 3, 0)
def equalizer(self, db: typing.Union[torch.Tensor, np.ndarray]):
"""Applies a mel-spaced equalizer to the audio signal.
Parameters
----------
db : typing.Union[torch.Tensor, np.ndarray]
EQ curve to apply.
Returns
-------
AudioSignal
AudioSignal with equalization applied.
"""
db = util.ensure_tensor(db)
n_bands = db.shape[-1]
fbank = self.mel_filterbank(n_bands)
# If there's a batch dimension, make sure it's the same.
if db.ndim == 2:
if db.shape[0] != 1:
assert db.shape[0] == fbank.shape[0]
else:
db = db.unsqueeze(0)
weights = (10**db).to(self.device).float()
fbank = fbank * weights[:, None, None, :]
eq_audio_data = fbank.sum(-1)
self.audio_data = eq_audio_data
return self
def clip_distortion(
self, clip_percentile: typing.Union[torch.Tensor, np.ndarray, float]
):
"""Clips the signal at a given percentile. The higher it is,
the lower the threshold for clipping.
Parameters
----------
clip_percentile : typing.Union[torch.Tensor, np.ndarray, float]
Values are between 0.0 to 1.0. Typical values are 0.1 or below.
Returns
-------
AudioSignal
Audio signal with clipped audio data.
"""
clip_percentile = util.ensure_tensor(clip_percentile, ndim=1)
min_thresh = torch.quantile(self.audio_data, clip_percentile / 2, dim=-1)
max_thresh = torch.quantile(self.audio_data, 1 - (clip_percentile / 2), dim=-1)
nc = self.audio_data.shape[1]
min_thresh = min_thresh[:, :nc, :]
max_thresh = max_thresh[:, :nc, :]
self.audio_data = self.audio_data.clamp(min_thresh, max_thresh)
return self
def quantization(
self, quantization_channels: typing.Union[torch.Tensor, np.ndarray, int]
):
"""Applies quantization to the input waveform.
Parameters
----------
quantization_channels : typing.Union[torch.Tensor, np.ndarray, int]
Number of evenly spaced quantization channels to quantize
to.
Returns
-------
AudioSignal
Quantized AudioSignal.
"""
quantization_channels = util.ensure_tensor(quantization_channels, ndim=3)
x = self.audio_data
x = (x + 1) / 2
x = x * quantization_channels
x = x.floor()
x = x / quantization_channels
x = 2 * x - 1
residual = (self.audio_data - x).detach()
self.audio_data = self.audio_data - residual
return self
def mulaw_quantization(
self, quantization_channels: typing.Union[torch.Tensor, np.ndarray, int]
):
"""Applies mu-law quantization to the input waveform.
Parameters
----------
quantization_channels : typing.Union[torch.Tensor, np.ndarray, int]
Number of mu-law spaced quantization channels to quantize
to.
Returns
-------
AudioSignal
Quantized AudioSignal.
"""
mu = quantization_channels - 1.0
mu = util.ensure_tensor(mu, ndim=3)
x = self.audio_data
# quantize
x = torch.sign(x) * torch.log1p(mu * torch.abs(x)) / torch.log1p(mu)
x = ((x + 1) / 2 * mu + 0.5).to(torch.int64)
# unquantize
x = (x / mu) * 2 - 1.0
x = torch.sign(x) * (torch.exp(torch.abs(x) * torch.log1p(mu)) - 1.0) / mu
residual = (self.audio_data - x).detach()
self.audio_data = self.audio_data - residual
return self
def __matmul__(self, other):
return self.convolve(other)
class ImpulseResponseMixin:
"""These functions are generally only used with AudioSignals that are derived
from impulse responses, not other sources like music or speech. These methods
are used to replicate the data augmentation described in [1].
1. Bryan, Nicholas J. "Impulse response data augmentation and deep
neural networks for blind room acoustic parameter estimation."
ICASSP 2020-2020 IEEE International Conference on Acoustics,
Speech and Signal Processing (ICASSP). IEEE, 2020.
"""
def decompose_ir(self):
"""Decomposes an impulse response into early and late
field responses.
"""
# Equations 1 and 2
# -----------------
# Breaking up into early
# response + late field response.
td = torch.argmax(self.audio_data, dim=-1, keepdim=True)
t0 = int(self.sample_rate * 0.0025)
idx = torch.arange(self.audio_data.shape[-1], device=self.device)[None, None, :]
idx = idx.expand(self.batch_size, -1, -1)
early_idx = (idx >= td - t0) * (idx <= td + t0)
early_response = torch.zeros_like(self.audio_data, device=self.device)
early_response[early_idx] = self.audio_data[early_idx]
late_idx = ~early_idx
late_field = torch.zeros_like(self.audio_data, device=self.device)
late_field[late_idx] = self.audio_data[late_idx]
# Equation 4
# ----------
# Decompose early response into windowed
# direct path and windowed residual.
window = torch.zeros_like(self.audio_data, device=self.device)
for idx in range(self.batch_size):
window_idx = early_idx[idx, 0].nonzero()
window[idx, ..., window_idx] = self.get_window(
"hann", window_idx.shape[-1], self.device
)
return early_response, late_field, window
def measure_drr(self):
"""Measures the direct-to-reverberant ratio of the impulse
response.
Returns
-------
float
Direct-to-reverberant ratio
"""
early_response, late_field, _ = self.decompose_ir()
num = (early_response**2).sum(dim=-1)
den = (late_field**2).sum(dim=-1)
drr = 10 * torch.log10(num / den)
return drr
@staticmethod
def solve_alpha(early_response, late_field, wd, target_drr):
"""Used to solve for the alpha value, which is used
to alter the drr.
"""
# Equation 5
# ----------
# Apply the good ol' quadratic formula.
wd_sq = wd**2
wd_sq_1 = (1 - wd) ** 2
e_sq = early_response**2
l_sq = late_field**2
a = (wd_sq * e_sq).sum(dim=-1)
b = (2 * (1 - wd) * wd * e_sq).sum(dim=-1)
c = (wd_sq_1 * e_sq).sum(dim=-1) - torch.pow(10, target_drr / 10) * l_sq.sum(
dim=-1
)
expr = ((b**2) - 4 * a * c).sqrt()
alpha = torch.maximum(
(-b - expr) / (2 * a),
(-b + expr) / (2 * a),
)
return alpha
def alter_drr(self, drr: typing.Union[torch.Tensor, np.ndarray, float]):
"""Alters the direct-to-reverberant ratio of the impulse response.
Parameters
----------
drr : typing.Union[torch.Tensor, np.ndarray, float]
Direct-to-reverberant ratio that impulse response will be
altered to, if specified, by default None
Returns
-------
AudioSignal
Altered impulse response.
"""
drr = util.ensure_tensor(drr, 2, self.batch_size).to(self.device)
early_response, late_field, window = self.decompose_ir()
alpha = self.solve_alpha(early_response, late_field, window, drr)
min_alpha = (
late_field.abs().max(dim=-1)[0] / early_response.abs().max(dim=-1)[0]
)
alpha = torch.maximum(alpha, min_alpha)[..., None]
aug_ir_data = (
alpha * window * early_response
+ ((1 - window) * early_response)
+ late_field
)
self.audio_data = aug_ir_data
self.ensure_max_of_audio()
return self