import torch import torch.nn as nn import numpy as np import typing import scipy from typing import List, Any from einops import rearrange from collections import namedtuple from torch.nn import functional as F from torchaudio.transforms import MelSpectrogram # GAN Loss class LSGANLoss(nn.Module): """ Computes a discriminator loss, given a discriminator on generated waveforms/spectrograms compared to ground truth waveforms/spectrograms. Computes the loss for both the discriminator and the generator in separate functions. """ def __init__(self, discriminator): super().__init__() self.discriminator = discriminator def forward(self, fake, real): d_fake = self.discriminator(fake) d_real = self.discriminator(real) return d_fake, d_real def discriminator_loss(self, fake, real): d_fake, d_real = self.forward(fake.clone().detach(), real) # loss_d = 0 # for x_fake, x_real in zip(d_fake, d_real): # loss_d += torch.mean(x_fake[-1] ** 2) # loss_d += torch.mean((1 - x_real[-1]) ** 2) # label smoothing loss_d = 0 real_label = 0.9 fake_label = 0.0 for x_fake, x_real in zip(d_fake, d_real): loss_d += torch.mean((x_fake[-1] - fake_label) ** 2) # Target: 0 loss_d += torch.mean((x_real[-1] - real_label) ** 2) # Target: 0.9 (not 1.0!) return loss_d def generator_loss(self, fake, real): d_fake, d_real = self.forward(fake, real) loss_g = 0 for x_fake in d_fake: loss_g += torch.mean((1 - x_fake[-1]) ** 2) loss_feature = 0 for i in range(len(d_fake)): for j in range(len(d_fake[i]) - 1): loss_feature += F.l1_loss(d_fake[i][j], d_real[i][j].detach()) return loss_g, loss_feature # Hinge loss class HingeLoss(nn.Module): def __init__(self, discriminator, normalize=True): super().__init__() self.discriminator = discriminator self.normalize = normalize def forward(self, fake, real): d_fake, fmap_fake = self.discriminator(fake) d_real, fmap_real = self.discriminator(real) return (d_fake, fmap_fake), (d_real, fmap_real) @staticmethod def hinge_loss(x): if x.numel() == 0: return torch.tensor([0.0], device=x.device) return F.relu(1 - x).mean() # from hilcodec # return -x.mean() @staticmethod def hinge_real_loss(x): return F.relu(1 - x).mean() # from hilcodec # return -torch.mean(torch.min(x - 1, torch.tensor(0., device=x.device).expand_as(x))) @staticmethod def hinge_fake_loss(x): return F.relu(1 + x).mean() # from hilcodec # return -torch.mean(torch.min(-x - 1, torch.tensor(0., device=x.device).expand_as(x))) def discriminator_loss(self, fake, real): loss = torch.tensor(0.0, device=fake.device) # Detach both fake and real inputs d_fake, _ = self.discriminator(fake.detach()) d_real, _ = self.discriminator(real.detach()) n_sub_discriminators = len(d_fake) for fake_pred, real_pred in zip(d_fake, d_real): loss += self.hinge_fake_loss(fake_pred) + self.hinge_real_loss(real_pred) if self.normalize: loss /= n_sub_discriminators return loss def generator_loss(self, fake, real): adv = torch.tensor(0.0, device=fake.device) feat = torch.tensor(0.0, device=fake.device) d_fake, fmap_fake = self.discriminator(fake) d_real, fmap_real = self.discriminator(real) n_sub_discriminators = len(d_fake) for fake_pred in d_fake: # push the fake pred to get closer to 1 adv += self.hinge_loss(fake_pred) # Feature matching loss for i in range(len(fmap_fake)): for j in range(len(fmap_fake[i])): feat += F.l1_loss(fmap_fake[i][j], fmap_real[i][j]) if self.normalize: adv /= n_sub_discriminators feat /= n_sub_discriminators return adv, feat # depricated class DepHingeLoss(nn.Module): def __init__(self, discriminator): super().__init__() self.discriminator = discriminator def forward(self, fake, real): d_fake, fmap_fake = self.discriminator(fake) d_real, fmap_real = self.discriminator(real) return (d_fake, fmap_fake), (d_real, fmap_real) def discriminator_loss(self, fake, real): (d_fake, _), (d_real, _) = self.forward(fake.clone().detach(), real) loss_d = 0 for fake_pred, real_pred in zip(d_fake, d_real): loss_d += torch.mean(torch.relu(1 + fake_pred[-1])) + torch.mean( torch.relu(1 - real_pred[-1]) ) return loss_d def generator_loss(self, fake, real): (d_fake, fmap_fake), (d_real, fmap_real) = self.forward(fake, real) # Generator adversarial loss loss_g = 0 for fake_pred in d_fake: loss_g += -torch.mean(fake_pred[-1]) # Feature matching loss loss_feature = 0 for i in range(len(fmap_fake)): for j in range(len(fmap_fake[i])): loss_feature += F.l1_loss(fmap_fake[i][j], fmap_real[i][j].detach()) return loss_g, loss_feature # Multi-scale melspectrogram loss STFTParams = namedtuple( "STFTParams", ["window_length", "hop_length", "window_type", "match_stride", "padding_type"], ) """ STFTParams object is a container that holds STFT parameters - window_length, hop_length, and window_type. Not all parameters need to be specified. Ones that are not specified will be inferred by the AudioSignal parameters. Parameters ---------- window_length : int, optional Window length of STFT, by default ``0.032 * self.sample_rate``. hop_length : int, optional Hop length of STFT, by default ``window_length // 4``. window_type : str, optional Type of window to use, by default ``sqrt\_hann``. match_stride : bool, optional Whether to match the stride of convolutional layers, by default False padding_type : str, optional Type of padding to use, by default 'reflect' """ STFTParams.__new__.__defaults__ = (None, None, None, None, None) def apply_reduction(losses, reduction="none"): """Apply reduction to collection of losses.""" if reduction == "mean": losses = losses.mean() elif reduction == "sum": losses = losses.sum() return losses def get_window(win_type: str, win_length: int): """Return a window function. Args: win_type (str): Window type. Can either be one of the window function provided in PyTorch ['hann_window', 'bartlett_window', 'blackman_window', 'hamming_window', 'kaiser_window'] or any of the windows provided by [SciPy](https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.windows.get_window.html). win_length (int): Window length Returns: win: The window as a 1D torch tensor """ try: win = getattr(torch, win_type)(win_length) except: win = torch.from_numpy(scipy.signal.windows.get_window(win_type, win_length)) return win class FIRFilter(torch.nn.Module): """FIR pre-emphasis filtering module. Args: filter_type (str): Shape of the desired FIR filter ("hp", "fd", "aw"). Default: "hp" coef (float): Coefficient value for the filter tap (only applicable for "hp" and "fd"). Default: 0.85 ntaps (int): Number of FIR filter taps for constructing A-weighting filters. Default: 101 plot (bool): Plot the magnitude respond of the filter. Default: False Based upon the perceptual loss pre-empahsis filters proposed by [Wright & Välimäki, 2019](https://arxiv.org/abs/1911.08922). A-weighting filter - "aw" First-order highpass - "hp" Folded differentiator - "fd" Note that the default coefficeint value of 0.85 is optimized for a sampling rate of 44.1 kHz, considering adjusting this value at differnt sampling rates. """ def __init__(self, filter_type="hp", coef=0.85, fs=44100, ntaps=101, plot=False): """Initilize FIR pre-emphasis filtering module.""" super(FIRFilter, self).__init__() self.filter_type = filter_type self.coef = coef self.fs = fs self.ntaps = ntaps self.plot = plot import scipy.signal if ntaps % 2 == 0: raise ValueError(f"ntaps must be odd (ntaps={ntaps}).") if filter_type == "hp": self.fir = torch.nn.Conv1d(1, 1, kernel_size=3, bias=False, padding=1) self.fir.weight.requires_grad = False self.fir.weight.data = torch.tensor([1, -coef, 0]).view(1, 1, -1) elif filter_type == "fd": self.fir = torch.nn.Conv1d(1, 1, kernel_size=3, bias=False, padding=1) self.fir.weight.requires_grad = False self.fir.weight.data = torch.tensor([1, 0, -coef]).view(1, 1, -1) elif filter_type == "aw": # Definition of analog A-weighting filter according to IEC/CD 1672. f1 = 20.598997 f2 = 107.65265 f3 = 737.86223 f4 = 12194.217 A1000 = 1.9997 NUMs = [(2 * np.pi * f4) ** 2 * (10 ** (A1000 / 20)), 0, 0, 0, 0] DENs = np.polymul( [1, 4 * np.pi * f4, (2 * np.pi * f4) ** 2], [1, 4 * np.pi * f1, (2 * np.pi * f1) ** 2], ) DENs = np.polymul(np.polymul(DENs, [1, 2 * np.pi * f3]), [1, 2 * np.pi * f2]) # convert analog filter to digital filter b, a = scipy.signal.bilinear(NUMs, DENs, fs=fs) # compute the digital filter frequency response w_iir, h_iir = scipy.signal.freqz(b, a, worN=512, fs=fs) # then we fit to 101 tap FIR filter with least squares taps = scipy.signal.firls(ntaps, w_iir, abs(h_iir), fs=fs) # now implement this digital FIR filter as a Conv1d layer self.fir = torch.nn.Conv1d(1, 1, kernel_size=ntaps, bias=False, padding=ntaps // 2) self.fir.weight.requires_grad = False self.fir.weight.data = torch.tensor(taps.astype("float32")).view(1, 1, -1) def forward(self, input, target): """Calculate forward propagation. Args: input (Tensor): Predicted signal (B, #channels, #samples). target (Tensor): Groundtruth signal (B, #channels, #samples). Returns: Tensor: Filtered signal. """ input = torch.nn.functional.conv1d(input, self.fir.weight.data, padding=self.ntaps // 2) target = torch.nn.functional.conv1d(target, self.fir.weight.data, padding=self.ntaps // 2) return input, target class SpectralConvergenceLoss(torch.nn.Module): """Spectral convergence loss module. See [Arik et al., 2018](https://arxiv.org/abs/1808.06719). """ def __init__(self): super(SpectralConvergenceLoss, self).__init__() def forward(self, x_mag, y_mag): return ( torch.norm(y_mag - x_mag, p="fro", dim=[-1, -2]) / torch.norm(y_mag, p="fro", dim=[-1, -2]) ).mean() class STFTMagnitudeLoss(torch.nn.Module): """STFT magnitude loss module. See [Arik et al., 2018](https://arxiv.org/abs/1808.06719) and [Engel et al., 2020](https://arxiv.org/abs/2001.04643v1) Log-magnitudes are calculated with `log(log_fac*x + log_eps)`, where `log_fac` controls the compression strength (larger value results in more compression), and `log_eps` can be used to control the range of the compressed output values (e.g., `log_eps>=1` ensures positive output values). The default values `log_fac=1` and `log_eps=0` correspond to plain log-compression. Args: log (bool, optional): Log-scale the STFT magnitudes, or use linear scale. Default: True log_eps (float, optional): Constant value added to the magnitudes before evaluating the logarithm. Default: 0.0 log_fac (float, optional): Constant multiplication factor for the magnitudes before evaluating the logarithm. Default: 1.0 distance (str, optional): Distance function ["L1", "L2"]. Default: "L1" reduction (str, optional): Reduction of the loss elements. Default: "mean" """ def __init__(self, log=True, log_eps=0.0, log_fac=1.0, distance="L1", reduction="mean"): super(STFTMagnitudeLoss, self).__init__() self.log = log self.log_eps = log_eps self.log_fac = log_fac if distance == "L1": self.distance = torch.nn.L1Loss(reduction=reduction) elif distance == "L2": self.distance = torch.nn.MSELoss(reduction=reduction) else: raise ValueError(f"Invalid distance: '{distance}'.") def forward(self, x_mag, y_mag): if self.log: x_mag = torch.log(self.log_fac * x_mag + self.log_eps) y_mag = torch.log(self.log_fac * y_mag + self.log_eps) return self.distance(x_mag, y_mag) class STFTLoss(torch.nn.Module): """STFT loss module. See [Yamamoto et al. 2019](https://arxiv.org/abs/1904.04472). Args: fft_size (int, optional): FFT size in samples. Default: 1024 hop_size (int, optional): Hop size of the FFT in samples. Default: 256 win_length (int, optional): Length of the FFT analysis window. Default: 1024 window (str, optional): Window to apply before FFT, can either be one of the window function provided in PyTorch ['hann_window', 'bartlett_window', 'blackman_window', 'hamming_window', 'kaiser_window'] or any of the windows provided by [SciPy](https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.windows.get_window.html). Default: 'hann_window' w_sc (float, optional): Weight of the spectral convergence loss term. Default: 1.0 w_log_mag (float, optional): Weight of the log magnitude loss term. Default: 1.0 w_lin_mag_mag (float, optional): Weight of the linear magnitude loss term. Default: 0.0 w_phs (float, optional): Weight of the spectral phase loss term. Default: 0.0 sample_rate (int, optional): Sample rate. Required when scale = 'mel'. Default: None scale (str, optional): Optional frequency scaling method, options include: ['mel', 'chroma'] Default: None n_bins (int, optional): Number of scaling frequency bins. Default: None. perceptual_weighting (bool, optional): Apply perceptual A-weighting (Sample rate must be supplied). Default: False scale_invariance (bool, optional): Perform an optimal scaling of the target. Default: False eps (float, optional): Small epsilon value for stablity. Default: 1e-8 output (str, optional): Format of the loss returned. 'loss' : Return only the raw, aggregate loss term. 'full' : Return the raw loss, plus intermediate loss terms. Default: 'loss' reduction (str, optional): Specifies the reduction to apply to the output: 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Default: 'mean' mag_distance (str, optional): Distance function ["L1", "L2"] for the magnitude loss terms. device (str, optional): Place the filterbanks on specified device. Default: None Returns: loss: Aggreate loss term. Only returned if output='loss'. By default. loss, sc_mag_loss, log_mag_loss, lin_mag_loss, phs_loss: Aggregate and intermediate loss terms. Only returned if output='full'. """ def __init__( self, fft_size: int = 1024, hop_size: int = 256, win_length: int = 1024, window: str = "hann_window", w_sc: float = 1.0, w_log_mag: float = 1.0, w_lin_mag: float = 0.0, w_phs: float = 0.0, sample_rate: float = None, scale: str = None, n_bins: int = None, perceptual_weighting: bool = False, scale_invariance: bool = False, eps: float = 1e-8, output: str = "loss", reduction: str = "mean", mag_distance: str = "L1", device: Any = None, **kwargs, ): super().__init__() self.fft_size = fft_size self.hop_size = hop_size self.win_length = win_length self.window = get_window(window, win_length) self.w_sc = w_sc self.w_log_mag = w_log_mag self.w_lin_mag = w_lin_mag self.w_phs = w_phs self.sample_rate = sample_rate self.scale = scale self.n_bins = n_bins self.perceptual_weighting = perceptual_weighting self.scale_invariance = scale_invariance self.eps = eps self.output = output self.reduction = reduction self.mag_distance = mag_distance self.device = device self.phs_used = bool(self.w_phs) self.spectralconv = SpectralConvergenceLoss() self.logstft = STFTMagnitudeLoss(log=True, reduction=reduction, distance=mag_distance, **kwargs) self.linstft = STFTMagnitudeLoss(log=False, reduction=reduction, distance=mag_distance, **kwargs) # setup mel filterbank if scale is not None: try: import librosa.filters except Exception as e: print(e) print("Try `pip install auraloss[all]`.") if self.scale == "mel": assert sample_rate is not None # Must set sample rate to use mel scale assert n_bins <= fft_size # Must be more FFT bins than Mel bins fb = librosa.filters.mel(sr=sample_rate, n_fft=fft_size, n_mels=n_bins) fb = torch.tensor(fb).unsqueeze(0) elif self.scale == "chroma": assert sample_rate is not None # Must set sample rate to use chroma scale assert n_bins <= fft_size # Must be more FFT bins than chroma bins fb = librosa.filters.chroma(sr=sample_rate, n_fft=fft_size, n_chroma=n_bins) else: raise ValueError(f"Invalid scale: {self.scale}. Must be 'mel' or 'chroma'.") self.register_buffer("fb", fb) if scale is not None and device is not None: self.fb = self.fb.to(self.device) # move filterbank to device if self.perceptual_weighting: if sample_rate is None: raise ValueError("`sample_rate` must be supplied when `perceptual_weighting = True`.") self.prefilter = FIRFilter(filter_type="aw", fs=sample_rate) def stft(self, x): """Perform STFT. Args: x (Tensor): Input signal tensor (B, T). Returns: Tensor: x_mag, x_phs Magnitude and phase spectra (B, fft_size // 2 + 1, frames). """ x_stft = torch.stft( x, self.fft_size, self.hop_size, self.win_length, self.window, return_complex=True, ) x_mag = torch.sqrt(torch.clamp((x_stft.real**2) + (x_stft.imag**2), min=self.eps)) # torch.angle is expensive, so it is only evaluated if the values are used in the loss if self.phs_used: x_phs = torch.angle(x_stft) else: x_phs = None return x_mag, x_phs def forward(self, input: torch.Tensor, target: torch.Tensor): bs, chs, seq_len = input.size() if self.perceptual_weighting: # apply optional A-weighting via FIR filter # since FIRFilter only support mono audio we will move channels to batch dim input = input.view(bs * chs, 1, -1) target = target.view(bs * chs, 1, -1) # now apply the filter to both self.prefilter.to(input.device) input, target = self.prefilter(input, target) # now move the channels back input = input.view(bs, chs, -1) target = target.view(bs, chs, -1) # compute the magnitude and phase spectra of input and target self.window = self.window.to(input.device) x_mag, x_phs = self.stft(input.view(-1, input.size(-1))) y_mag, y_phs = self.stft(target.view(-1, target.size(-1))) # apply relevant transforms if self.scale is not None: self.fb = self.fb.to(input.device) x_mag = torch.matmul(self.fb, x_mag) y_mag = torch.matmul(self.fb, y_mag) # normalize scales if self.scale_invariance: alpha = (x_mag * y_mag).sum([-2, -1]) / ((y_mag**2).sum([-2, -1])) y_mag = y_mag * alpha.unsqueeze(-1) # compute loss terms sc_mag_loss = self.spectralconv(x_mag, y_mag) if self.w_sc else 0.0 log_mag_loss = self.logstft(x_mag, y_mag) if self.w_log_mag else 0.0 lin_mag_loss = self.linstft(x_mag, y_mag) if self.w_lin_mag else 0.0 phs_loss = torch.nn.functional.mse_loss(x_phs, y_phs) if self.phs_used else 0.0 # combine loss terms loss = ( (self.w_sc * sc_mag_loss) + (self.w_log_mag * log_mag_loss) + (self.w_lin_mag * lin_mag_loss) + (self.w_phs * phs_loss) ) loss = apply_reduction(loss, reduction=self.reduction) if self.output == "loss": return loss elif self.output == "full": return loss, sc_mag_loss, log_mag_loss, lin_mag_loss, phs_loss class MultiResolutionSTFTLoss(torch.nn.Module): """Multi resolution STFT loss module. See [Yamamoto et al., 2019](https://arxiv.org/abs/1910.11480) Args: fft_sizes (list): List of FFT sizes. hop_sizes (list): List of hop sizes. win_lengths (list): List of window lengths. window (str, optional): Window to apply before FFT, options include: 'hann_window', 'bartlett_window', 'blackman_window', 'hamming_window', 'kaiser_window'] Default: 'hann_window' w_sc (float, optional): Weight of the spectral convergence loss term. Default: 1.0 w_log_mag (float, optional): Weight of the log magnitude loss term. Default: 1.0 w_lin_mag (float, optional): Weight of the linear magnitude loss term. Default: 0.0 w_phs (float, optional): Weight of the spectral phase loss term. Default: 0.0 sample_rate (int, optional): Sample rate. Required when scale = 'mel'. Default: None scale (str, optional): Optional frequency scaling method, options include: ['mel', 'chroma'] Default: None n_bins (int, optional): Number of mel frequency bins. Required when scale = 'mel'. Default: None. scale_invariance (bool, optional): Perform an optimal scaling of the target. Default: False """ def __init__( self, fft_sizes: List[int] = [1024, 2048, 512], hop_sizes: List[int] = [120, 240, 50], win_lengths: List[int] = [600, 1200, 240], window: str = "hann_window", w_sc: float = 1.0, w_log_mag: float = 1.0, w_lin_mag: float = 0.0, w_phs: float = 0.0, sample_rate: float = None, scale: str = None, n_bins: int = None, perceptual_weighting: bool = False, scale_invariance: bool = False, **kwargs, ): super().__init__() assert len(fft_sizes) == len(hop_sizes) == len(win_lengths) # must define all self.fft_sizes = fft_sizes self.hop_sizes = hop_sizes self.win_lengths = win_lengths self.stft_losses = torch.nn.ModuleList() for fs, ss, wl in zip(fft_sizes, hop_sizes, win_lengths): self.stft_losses += [ STFTLoss( fs, ss, wl, window, w_sc, w_log_mag, w_lin_mag, w_phs, sample_rate, scale, n_bins, perceptual_weighting, scale_invariance, **kwargs, ) ] def forward(self, x, y): mrstft_loss = 0.0 sc_mag_loss, log_mag_loss, lin_mag_loss, phs_loss = [], [], [], [] for f in self.stft_losses: if f.output == "full": # extract just first term tmp_loss = f(x, y) mrstft_loss += tmp_loss[0] sc_mag_loss.append(tmp_loss[1]) log_mag_loss.append(tmp_loss[2]) lin_mag_loss.append(tmp_loss[3]) phs_loss.append(tmp_loss[4]) else: mrstft_loss += f(x, y) mrstft_loss /= len(self.stft_losses) if f.output == "loss": return mrstft_loss else: return mrstft_loss, sc_mag_loss, log_mag_loss, lin_mag_loss, phs_loss class MelSpectrogramLoss(nn.Module): """Compute distance between mel spectrograms. Can be used in a multi-scale way. Parameters ---------- n_mels : List[int] Number of mels per STFT, by default [150, 80], window_lengths : List[int], optional Length of each window of each STFT, by default [2048, 512] loss_fn : typing.Callable, optional How to compare each loss, by default nn.L1Loss() clamp_eps : float, optional Clamp on the log magnitude, below, by default 1e-5 mag_weight : float, optional Weight of raw magnitude portion of loss, by default 1.0 log_weight : float, optional Weight of log magnitude portion of loss, by default 1.0 pow : float, optional Power to raise magnitude to before taking log, by default 2.0 weight : float, optional Weight of this loss, by default 1.0 match_stride : bool, optional Whether to match the stride of convolutional layers, by default False Implementation copied from: https://github.com/descriptinc/lyrebird-audiotools/blob/961786aa1a9d628cca0c0486e5885a457fe70c1a/audiotools/metrics/spectral.py """ def __init__( self, n_mels: List[int] = [5, 10, 20, 40, 80, 160, 320], window_lengths: List[int] = [32, 64, 128, 256, 512, 1024, 2048], loss_fn: typing.Callable = nn.L1Loss(), clamp_eps: float = 1e-5, mag_weight: float = 1.0, log_weight: float = 1.0, pow: float = 1.0, weight: float = 1.0, match_stride: bool = False, mel_fmin: List[float] = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], mel_fmax: List[float] = [None, None, None, None, None, None, None], window_type: str = None, ): super().__init__() stft_params = [ STFTParams( window_length=w, hop_length=w // 4, match_stride=match_stride, window_type=window_type, ) for w in window_lengths ] self.n_mels = n_mels self.loss_fn = loss_fn self.clamp_eps = clamp_eps self.log_weight = log_weight self.mag_weight = mag_weight self.weight = weight self.mel_fmin = mel_fmin self.mel_fmax = mel_fmax self.pow = pow self.melspecs = nn.ModuleList( [ MelSpectrogram( sample_rate=48000, n_mels=n_mel, f_min=fmin, f_max=fmax, n_fft=s.window_length, win_length=s.window_length, hop_length=s.hop_length, power=pow, ) for n_mel, fmin, fmax, s in zip(n_mels, mel_fmin, mel_fmax, stft_params) ] ) self.eps = 1e-10 def forward(self, x, y): """Computes mel loss between an estimate and a reference signal. Parameters ---------- x : torch.tensor Estimate signal y : torch.tensor Reference signal Returns ------- torch.Tensor Mel loss. """ loss = 0.0 for mel_spec in self.melspecs: x_mels = mel_spec(x) y_mels = mel_spec(y) loss += self.loss_fn(torch.log10(x_mels + self.eps), torch.log10(y_mels + self.eps)) return loss # PEAQ loss def safe_pow(x, y): if x == 0 and y <= 0: if y == 0: return 1 else: return float("inf") return x**y def omem(f): return float( -2.184 * safe_pow(f / 1000, -0.8) + 6.5 * np.exp(-0.6 * safe_pow(f / 1000 - 3.3, 2)) - 0.001 * safe_pow(f / 1000, 3.6) ) class PEAQL1Loss(nn.Module): """Compute distance between Bark-scale spectrograms described in PEAQ. Parameters ---------- """ def __init__( self, sample_rate: int = 48000, window_length: int = 2048, overlap: int = 1024, loss_fn: typing.Callable = nn.L1Loss(), ): super().__init__() # hyper parameters self.sample_rate = sample_rate self.window_length = window_length self.overlap = overlap self.loss_fn = loss_fn # fixed parameters L_p = 92 A_max = 1 gamma_fc = 0.8497 G_L = 10 ** (L_p / 20) / (gamma_fc * A_max / 4 * (window_length - 1)) hann = torch.hann_window(window_length).numpy() weighted_hann = G_L * hann self.weighted_hann = weighted_hann[np.newaxis, :] self.omem_weights = self.get_omem() self.filterbank, self.f_c = self.get_filterbank() self.internal_noise = self.get_noise() def get_omem(self): """Outher and middle ear modeling weights""" center_freq = np.linspace(0, self.sample_rate // 2, self.window_length // 2 + 1).astype( np.float32 ) a_db = [omem(f) for f in center_freq] weights = torch.tensor([10 ** (a / 20) for a in a_db]).unsqueeze(0) return weights def get_filterbank(self): """Bark-scale filterbank""" f_l = np.array( [ 80.000, 103.445, 127.023, 150.762, 174.694, 198.849, 223.257, 247.950, 272.959, 298.317, 324.055, 350.207, 376.805, 403.884, 431.478, 459.622, 488.353, 517.707, 547.721, 578.434, 609.885, 642.114, 675.161, 709.071, 743.884, 779.647, 816.404, 854.203, 893.091, 933.119, 974.336, 1016.797, 1060.555, 1105.666, 1152.187, 1200.178, 1249.700, 1300.816, 1353.592, 1408.094, 1464.392, 1522.559, 1582.668, 1644.795, 1709.021, 1775.427, 1844.098, 1915.121, 1988.587, 2064.590, 2143.227, 2224.597, 2308.806, 2395.959, 2486.169, 2579.551, 2676.223, 2776.309, 2879.937, 2987.238, 3098.350, 3213.415, 3332.579, 3455.993, 3583.817, 3716.212, 3853.817, 3995.399, 4142.547, 4294.979, 4452.890, 4616.482, 4785.962, 4961.548, 5143.463, 5331.939, 5527.217, 5729.545, 5939.183, 6156.396, 6381.463, 6614.671, 6856.316, 7106.708, 7366.166, 7635.020, 7913.614, 8202.302, 8501.454, 8811.450, 9132.688, 9465.574, 9810.536, 10168.013, 10538.460, 10922.351, 11320.175, 11732.438, 12159.670, 12602.412, 13061.229, 13536.710, 14029.458, 14540.103, 15069.295, 15617.710, 16186.049, 16775.035, 17385.420, ] ) f_c = np.array( [ 91.708, 115.216, 138.870, 162.702, 186.742, 211.019, 235.566, 260.413, 285.593, 311.136, 337.077, 363.448, 390.282, 417.614, 445.479, 473.912, 502.950, 532.629, 562.988, 594.065, 625.899, 658.533, 692.006, 726.362, 761.644, 797.898, 835.170, 873.508, 912.959, 953.576, 995.408, 1038.511, 1082.938, 1128.746, 1175.995, 1224.744, 1275.055, 1326.992, 1380.623, 1436.014, 1493.237, 1552.366, 1613.474, 1676.641, 1741.946, 1809.474, 1879.310, 1951.543, 2026.266, 2103.573, 2183.564, 2266.340, 2352.008, 2440.675, 2532.456, 2627.468, 2725.832, 2827.672, 2933.120, 3042.309, 3155.379, 3272.475, 3393.745, 3519.344, 3649.432, 3784.176, 3923.748, 4068.324, 4218.090, 4373.237, 4533.963, 4700.473, 4872.978, 5051.700, 5236.866, 5428.712, 5627.484, 5833.434, 6046.825, 6267.931, 6497.031, 6734.420, 6980.399, 7235.284, 7499.397, 7773.077, 8056.673, 8350.547, 8655.072, 8970.639, 9297.648, 9636.520, 9987.683, 10351.586, 10728.695, 11119.490, 11524.470, 11944.149, 12379.066, 12829.775, 13294.850, 13780.887, 14282.503, 14802.338, 15341.057, 15899.345, 16477.914, 17077.504, 17690.045, ] ) f_u = np.array( [ 103.445, 127.023, 150.762, 174.694, 198.849, 223.257, 247.950, 272.959, 298.317, 324.055, 350.207, 376.805, 403.884, 431.478, 459.622, 488.353, 517.707, 547.721, 578.434, 609.885, 642.114, 675.161, 709.071, 743.884, 779.647, 816.404, 854.203, 893.091, 933.113, 974.336, 1016.797, 1060.555, 1105.666, 1152.187, 1200.178, 1249.700, 1300.816, 1353.592, 1408.094, 1464.392, 1522.559, 1582.668, 1644.795, 1709.021, 1775.427, 1844.098, 1915.121, 1988.587, 2064.590, 2143.227, 2224.597, 2308.806, 2395.959, 2486.169, 2579.551, 2676.223, 2776.309, 2879.937, 2987.238, 3098.350, 3213.415, 3332.579, 3455.993, 3583.817, 3716.212, 3853.348, 3995.399, 4142.547, 4294.979, 4452.890, 4643.482, 4785.962, 4961.548, 5143.463, 5331.939, 5527.217, 5729.545, 5939.183, 6156.396, 6381.463, 6614.671, 6856.316, 7106.708, 7366.166, 7635.020, 7913.614, 8202.302, 8501.454, 8811.450, 9132.688, 9465.574, 9810.536, 10168.013, 10538.460, 10922.351, 11320.175, 11732.438, 12159.670, 12602.412, 13061.229, 13536.710, 14029.458, 14540.103, 15069.295, 15617.710, 16186.049, 16775.035, 17385.420, 18000.000, ] ) df = self.sample_rate / self.window_length filterbank = torch.zeros(self.window_length // 2 + 1, len(f_c)) for k in range(filterbank.shape[0]): for i in range(filterbank.shape[1]): temp = (np.amin([f_u[i], (k + 0.5) * df]) - np.amax([f_l[i], (k - 0.5) * df])) / df filterbank[k, i] = np.amax([0, temp]) return filterbank, f_c def get_noise(self): """Internal ear noise""" return torch.tensor([10 ** (1.456 * ((f / 1000) ** -0.8) / 10) for f in self.f_c]).unsqueeze(0) def slicing(self, sequence: np.array): """ Input sequence: numpy array with a shape [batch, length] Output numpy array with a shape [batch, num_chunks, window_length] """ num_chunks = (sequence.shape[-1] - self.window_length) // (self.window_length - self.overlap) + 1 chunks = np.stack( [ np.lib.stride_tricks.sliding_window_view(s, self.window_length)[ :: self.window_length - self.overlap ] for s in sequence ] ) # batch, num_slice, window_length return chunks[:, :num_chunks, :] def windowing(self, stacked_wav): """Hann window""" return stacked_wav * self.weighted_hann def DFT(self, windowed_wav): """Discrete Fourier transform""" effective_bins = self.window_length // 2 spec = torch.fft.fft(torch.tensor(windowed_wav), n=self.window_length) squared_spec = torch.cat( [ spec[:, :1].real ** 2, spec[:, 1:effective_bins].real ** 2 + spec[:, 1:effective_bins].imag ** 2, spec[:, effective_bins : effective_bins + 1].real ** 2, ], dim=-1, ) return squared_spec def get_representation(self, wav): """ Input wav: torch.tensor [batch, 1, length] """ # input audio wav = wav.audio_data.squeeze(1).cpu().detach().numpy() b = len(wav) # slicing stacked_wav = self.slicing(wav) stacked_wav = rearrange(stacked_wav, "b c w -> (b c) w") # windowing windowed_wav = self.windowing(stacked_wav) # discrete Fourier transform squared_spec = self.DFT(windowed_wav) # outer and middle ear modeling weighted_spec = (self.omem_weights**2) * squared_spec # grouping filterbank grouped_spec = torch.matmul(weighted_spec, self.filterbank) grouped_spec[grouped_spec < 1e-12] = 1e-12 # internal ear noise noised_spec = grouped_spec + self.internal_noise # rearrange return rearrange(noised_spec, "(b c) f -> b c f", b=b) def forward(self, x, y): """Computes mel loss between an estimate and a reference signal. Parameters ---------- x : torch.tensor Estimate signal y : torch.tensor Reference signal Returns ------- torch.Tensor Mel loss. """ x_repr = self.get_representation(x) y_repr = self.get_representation(y) x_spec = 20 * torch.log10(x_repr) y_spec = 20 * torch.log10(y_repr) loss = self.loss_fn(x_spec, y_spec) return loss class PEAQSpreadL1Loss(nn.Module): """Compute distance between spreaded Bark-scale spectrograms described in PEAQ. Parameters ---------- """ def __init__( self, sample_rate: int = 48000, window_length: int = 2048, overlap: int = 1024, loss_fn: typing.Callable = nn.L1Loss(), ratio: float = 0.8, ): super().__init__() # hyper parameters self.sample_rate = sample_rate self.window_length = window_length self.overlap = overlap self.loss_fn = loss_fn self.ratio = ratio # fixed parameters L_p = 92 A_max = 1 gamma_fc = 0.8497 G_L = 10 ** (L_p / 20) / (gamma_fc * A_max / 4 * (window_length - 1)) hann = torch.hann_window(window_length).numpy() weighted_hann = G_L * hann self.weighted_hann = weighted_hann[np.newaxis, :] self.omem_weights = self.get_omem() self.filterbank, self.f_c = self.get_filterbank() self.internal_noise = self.get_noise() self.bs = self.frequency_spreading(torch.ones(1, len(self.f_c)), torch.ones(1, len(self.f_c))) self.alpha = self.get_alpha() def get_omem(self): """Outher and middle ear modeling weights""" center_freq = np.linspace(0, self.sample_rate // 2, self.window_length // 2 + 1).astype( np.float32 ) a_db = [omem(f) for f in center_freq] weights = torch.tensor([10 ** (a / 20) for a in a_db]).unsqueeze(0) return weights def get_filterbank(self): """Bark-scale filterbank""" f_l = np.array( [ 80.000, 103.445, 127.023, 150.762, 174.694, 198.849, 223.257, 247.950, 272.959, 298.317, 324.055, 350.207, 376.805, 403.884, 431.478, 459.622, 488.353, 517.707, 547.721, 578.434, 609.885, 642.114, 675.161, 709.071, 743.884, 779.647, 816.404, 854.203, 893.091, 933.119, 974.336, 1016.797, 1060.555, 1105.666, 1152.187, 1200.178, 1249.700, 1300.816, 1353.592, 1408.094, 1464.392, 1522.559, 1582.668, 1644.795, 1709.021, 1775.427, 1844.098, 1915.121, 1988.587, 2064.590, 2143.227, 2224.597, 2308.806, 2395.959, 2486.169, 2579.551, 2676.223, 2776.309, 2879.937, 2987.238, 3098.350, 3213.415, 3332.579, 3455.993, 3583.817, 3716.212, 3853.817, 3995.399, 4142.547, 4294.979, 4452.890, 4616.482, 4785.962, 4961.548, 5143.463, 5331.939, 5527.217, 5729.545, 5939.183, 6156.396, 6381.463, 6614.671, 6856.316, 7106.708, 7366.166, 7635.020, 7913.614, 8202.302, 8501.454, 8811.450, 9132.688, 9465.574, 9810.536, 10168.013, 10538.460, 10922.351, 11320.175, 11732.438, 12159.670, 12602.412, 13061.229, 13536.710, 14029.458, 14540.103, 15069.295, 15617.710, 16186.049, 16775.035, 17385.420, ] ) f_c = np.array( [ 91.708, 115.216, 138.870, 162.702, 186.742, 211.019, 235.566, 260.413, 285.593, 311.136, 337.077, 363.448, 390.282, 417.614, 445.479, 473.912, 502.950, 532.629, 562.988, 594.065, 625.899, 658.533, 692.006, 726.362, 761.644, 797.898, 835.170, 873.508, 912.959, 953.576, 995.408, 1038.511, 1082.938, 1128.746, 1175.995, 1224.744, 1275.055, 1326.992, 1380.623, 1436.014, 1493.237, 1552.366, 1613.474, 1676.641, 1741.946, 1809.474, 1879.310, 1951.543, 2026.266, 2103.573, 2183.564, 2266.340, 2352.008, 2440.675, 2532.456, 2627.468, 2725.832, 2827.672, 2933.120, 3042.309, 3155.379, 3272.475, 3393.745, 3519.344, 3649.432, 3784.176, 3923.748, 4068.324, 4218.090, 4373.237, 4533.963, 4700.473, 4872.978, 5051.700, 5236.866, 5428.712, 5627.484, 5833.434, 6046.825, 6267.931, 6497.031, 6734.420, 6980.399, 7235.284, 7499.397, 7773.077, 8056.673, 8350.547, 8655.072, 8970.639, 9297.648, 9636.520, 9987.683, 10351.586, 10728.695, 11119.490, 11524.470, 11944.149, 12379.066, 12829.775, 13294.850, 13780.887, 14282.503, 14802.338, 15341.057, 15899.345, 16477.914, 17077.504, 17690.045, ] ) f_u = np.array( [ 103.445, 127.023, 150.762, 174.694, 198.849, 223.257, 247.950, 272.959, 298.317, 324.055, 350.207, 376.805, 403.884, 431.478, 459.622, 488.353, 517.707, 547.721, 578.434, 609.885, 642.114, 675.161, 709.071, 743.884, 779.647, 816.404, 854.203, 893.091, 933.113, 974.336, 1016.797, 1060.555, 1105.666, 1152.187, 1200.178, 1249.700, 1300.816, 1353.592, 1408.094, 1464.392, 1522.559, 1582.668, 1644.795, 1709.021, 1775.427, 1844.098, 1915.121, 1988.587, 2064.590, 2143.227, 2224.597, 2308.806, 2395.959, 2486.169, 2579.551, 2676.223, 2776.309, 2879.937, 2987.238, 3098.350, 3213.415, 3332.579, 3455.993, 3583.817, 3716.212, 3853.348, 3995.399, 4142.547, 4294.979, 4452.890, 4643.482, 4785.962, 4961.548, 5143.463, 5331.939, 5527.217, 5729.545, 5939.183, 6156.396, 6381.463, 6614.671, 6856.316, 7106.708, 7366.166, 7635.020, 7913.614, 8202.302, 8501.454, 8811.450, 9132.688, 9465.574, 9810.536, 10168.013, 10538.460, 10922.351, 11320.175, 11732.438, 12159.670, 12602.412, 13061.229, 13536.710, 14029.458, 14540.103, 15069.295, 15617.710, 16186.049, 16775.035, 17385.420, 18000.000, ] ) df = self.sample_rate / self.window_length filterbank = torch.zeros(self.window_length // 2 + 1, len(f_c)) for k in range(filterbank.shape[0]): for i in range(filterbank.shape[1]): temp = (np.amin([f_u[i], (k + 0.5) * df]) - np.amax([f_l[i], (k - 0.5) * df])) / df filterbank[k, i] = np.amax([0, temp]) return filterbank, f_c def get_noise(self): """Internal ear noise""" return torch.tensor([10 ** (1.456 * ((f / 1000) ** -0.8) / 10) for f in self.f_c]).unsqueeze(0) def get_alpha(self): """Get alpha for time spreading""" tau_100 = 0.030 tau_min = 0.008 Nadv = self.window_length / 2 Fss = self.sample_rate / Nadv tau = tau_min + (np.divide(100, self.f_c)) * (tau_100 - tau_min) alpha = np.exp(np.divide(-1 / Fss, tau)) return torch.tensor(alpha).unsqueeze(0) def slicing(self, sequence: np.array): """ Input sequence: numpy array with a shape [batch, length] Output numpy array with a shape [batch, num_chunks, window_length] """ num_chunks = (sequence.shape[-1] - self.window_length) // (self.window_length - self.overlap) + 1 chunks = np.stack( [ np.lib.stride_tricks.sliding_window_view(s, self.window_length)[ :: self.window_length - self.overlap ] for s in sequence ] ) # batch, num_slice, window_length return chunks[:, :num_chunks, :] def windowing(self, stacked_wav): """Hann window""" return stacked_wav * self.weighted_hann def DFT(self, windowed_wav): """Discrete Fourier transform""" effective_bins = self.window_length // 2 spec = torch.fft.fft(torch.tensor(windowed_wav), n=self.window_length) squared_spec = torch.cat( [ spec[:, :1].real ** 2, spec[:, 1:effective_bins].real ** 2 + spec[:, 1:effective_bins].imag ** 2, spec[:, effective_bins : effective_bins + 1].real ** 2, ], dim=-1, ) return squared_spec def frequency_spreading(self, E, Bs): """frequency spreading""" e = 0.4 dz = 0.25 aL = 10 ** (-2.7 * dz) aUC = torch.tensor(10 ** ((-2.4 - 23 / self.f_c) * dz)) aUCE = aUC * (E ** (0.2 * dz)) gIL = torch.tensor([(1 - aL ** (i + 1)) / (1 - aL) for i in range(len(self.f_c))]).unsqueeze(0) gIU = torch.stack( [(1 - aUCE[:, i] ** (len(self.f_c) - i)) / (1 - aUCE[:, i]) for i in range(len(self.f_c))] ).T En = E / (gIL + gIU - 1) aUCEe = aUCE**e Ene = En**e # lower spreading Es = torch.zeros(Ene.shape) Es[:, -1] = Ene[:, -1] aLe = aL**e for i in range(len(self.f_c) - 2, -1, -1): Es[:, i] = aLe * Es[:, i + 1] + Ene[:, i] # upper spreading for i in range(len(self.f_c) - 1): r = Ene[:, i] a = aUCEe[:, i] for j in range(i + 1, len(self.f_c)): r = r * a Es[:, j] += r Es = (Es ** (1 / e)) / Bs return Es def time_spreading(self, Es): """time spreading""" b, t, f = Es.shape prev = torch.zeros(b, f) for i in range(t): prev = self.alpha * prev + (1 - self.alpha) * Es[:, i, :] Es[:, i, :] = torch.max(prev, Es[:, i, :]) return Es def adbb(self, E_test, E_ref): """average distorted block""" c = [-0.198719, 0.0550197, -0.00102438, 5.05622e-6, 9.01033e-11] d1 = 5.95072 d2 = 6.39468 g = 1.71332 bP = 4 bM = 6 E_test = 10 * np.log10(E_test.numpy()) E_ref = 10 * np.log10(E_ref.numpy()) # asymmetric excitation E_diff = E_ref - E_test L = np.zeros(E_test.shape) L[E_diff > 0] = 0.3 * E_ref[E_diff > 0] + 0.7 * E_test[E_diff > 0] L[E_diff <= 0] = E_test[E_diff <= 0] # effective detection step size (just noticeable difference) s = np.zeros(L.shape) s[L > 0] = ( c[0] + c[1] * L[L > 0] + c[2] * (L[L > 0] ** 2) + c[3] * (L[L > 0] ** 3) + c[4] * (L[L > 0] ** 4) + d1 * (d2 / L[L > 0]) ** g ) s[L <= 0] = 10**30 # detection probability above threshold p = np.zeros(s.shape) p[E_diff > 0] = 1 - 0.5 ** ((E_diff[E_diff > 0] / s[E_diff > 0]) ** bP) p[E_diff <= 0] = 1 - 0.5 ** ((E_diff[E_diff <= 0] / s[E_diff <= 0]) ** bM) q = np.abs(np.floor(E_diff)) / s # average counter = p[p > 0.5].sum() if counter == 0: adbb = 0.0 else: adbb = np.log10(q[p > 0.5].sum() / p[p > 0.5].sum()) return adbb def get_representation(self, wav): """ Input wav: torch.tensor [batch, 1, length] """ # input audio wav = wav.squeeze(1).cpu().detach().float().numpy() b = len(wav) # slicing stacked_wav = self.slicing(wav) stacked_wav = rearrange(stacked_wav, "b t f -> (b t) f") # windowing windowed_wav = self.windowing(stacked_wav) # discrete Fourier transform squared_spec = self.DFT(windowed_wav) # outer and middle ear modeling weighted_spec = (self.omem_weights**2) * squared_spec # grouping filterbank grouped_spec = torch.matmul(weighted_spec, self.filterbank) grouped_spec[grouped_spec < 1e-12] = 1e-12 # internal ear noise noised_spec = grouped_spec + self.internal_noise # frequency spreading freq_spread_spec = self.frequency_spreading(noised_spec, self.bs) # rearrange noised_repr = rearrange(noised_spec, "(b t) f -> b t f", b=b) freq_spread_repr = rearrange(freq_spread_spec, "(b t) f -> b t f", b=b) # time spreading time_spread_repr = self.time_spreading(freq_spread_repr) return ( 20 * torch.log10(noised_repr), 20 * torch.log10(freq_spread_repr), # 20 * torch.log10(time_spread_repr), time_spread_repr, ) def forward(self, x, y): """Computes mel loss between an estimate and a reference signal. Parameters ---------- x : torch.tensor Estimate signal y : torch.tensor Reference signal Returns ------- torch.Tensor Mel loss. """ # process x_repr_noised, x_repr_spread, E_test = self.get_representation(x) y_repr_noised, y_repr_spread, E_ref = self.get_representation(y) # adbb # adbb = self.adbb(E_test, E_ref) loss_noised = (1 - self.ratio) * self.loss_fn(x_repr_noised, y_repr_noised) loss_spread = self.ratio * self.loss_fn(x_repr_spread, y_repr_spread) return loss_noised + loss_spread