Spaces:
Running
on
Zero
Running
on
Zero
| import numpy as np | |
| import torch | |
| import librosa | |
| import torch.nn.functional as F | |
| from typing import Dict, List, Tuple | |
| def sdr(references: np.ndarray, estimates: np.ndarray) -> np.ndarray: | |
| """ | |
| Compute Signal-to-Distortion Ratio (SDR) for one or more audio tracks. | |
| SDR is a measure of how well the predicted source (estimate) matches the reference source. | |
| It is calculated as the ratio of the energy of the reference signal to the energy of the error (difference between reference and estimate). | |
| Return SDR in decibels (dB) | |
| Parameters: | |
| ---------- | |
| references : np.ndarray | |
| A 3D numpy array of shape (num_sources, num_channels, num_samples), where num_sources is the number of sources, | |
| num_channels is the number of channels (e.g., 1 for mono, 2 for stereo), and num_samples is the length of the audio signal. | |
| estimates : np.ndarray | |
| A 3D numpy array of shape (num_sources, num_channels, num_samples) representing the estimated sources. | |
| Returns: | |
| ------- | |
| np.ndarray | |
| A 1D numpy array containing the SDR values for each source. | |
| """ | |
| eps = 1e-8 # to avoid numerical errors | |
| num = np.sum(np.square(references), axis=(1, 2)) | |
| den = np.sum(np.square(references - estimates), axis=(1, 2)) | |
| num += eps | |
| den += eps | |
| return 10 * np.log10(num / den) | |
| def si_sdr(reference: np.ndarray, estimate: np.ndarray) -> float: | |
| """ | |
| Compute Scale-Invariant Signal-to-Distortion Ratio (SI-SDR) for one or more audio tracks. | |
| SI-SDR is a variant of the SDR metric that is invariant to the scaling of the estimate relative to the reference. | |
| It is calculated by scaling the estimate to match the reference signal and then computing the SDR. | |
| Parameters: | |
| ---------- | |
| reference : np.ndarray | |
| A 3D numpy array of shape (num_sources, num_channels, num_samples), where num_sources is the number of sources, | |
| num_channels is the number of channels (e.g., 1 for mono, 2 for stereo), and num_samples is the length of the audio signal. | |
| estimate : np.ndarray | |
| A 3D numpy array of shape (num_sources, num_channels, num_samples) representing the estimated sources. | |
| Returns: | |
| ------- | |
| float | |
| The SI-SDR value for the source. It is a scalar representing the Signal-to-Distortion Ratio in decibels (dB). | |
| """ | |
| eps = 1e-8 # To avoid numerical errors | |
| scale = np.sum(estimate * reference + eps, axis=(0, 1)) / np.sum(reference ** 2 + eps, axis=(0, 1)) | |
| scale = np.expand_dims(scale, axis=(0, 1)) # Reshape to [num_sources, 1] | |
| reference = reference * scale | |
| si_sdr = np.mean(10 * np.log10( | |
| np.sum(reference ** 2, axis=(0, 1)) / (np.sum((reference - estimate) ** 2, axis=(0, 1)) + eps) + eps)) | |
| return si_sdr | |
| def L1Freq_metric( | |
| reference: np.ndarray, | |
| estimate: np.ndarray, | |
| fft_size: int = 2048, | |
| hop_size: int = 1024, | |
| device: str = 'cpu' | |
| ) -> float: | |
| """ | |
| Compute the L1 Frequency Metric between the reference and estimated audio signals. | |
| This metric compares the magnitude spectrograms of the reference and estimated audio signals | |
| using the Short-Time Fourier Transform (STFT) and calculates the L1 loss between them. The result | |
| is scaled to the range [0, 100] where a higher value indicates better performance. | |
| Parameters: | |
| ---------- | |
| reference : np.ndarray | |
| A 2D numpy array of shape (num_channels, num_samples) representing the reference (ground truth) audio signal. | |
| estimate : np.ndarray | |
| A 2D numpy array of shape (num_channels, num_samples) representing the estimated (predicted) audio signal. | |
| fft_size : int, optional | |
| The size of the FFT (Short-Time Fourier Transform). Default is 2048. | |
| hop_size : int, optional | |
| The hop size between STFT frames. Default is 1024. | |
| device : str, optional | |
| The device to run the computation on ('cpu' or 'cuda'). Default is 'cpu'. | |
| Returns: | |
| ------- | |
| float | |
| The L1 Frequency Metric in the range [0, 100], where higher values indicate better performance. | |
| """ | |
| reference = torch.from_numpy(reference).to(device) | |
| estimate = torch.from_numpy(estimate).to(device) | |
| reference_stft = torch.stft(reference, fft_size, hop_size, return_complex=True) | |
| estimated_stft = torch.stft(estimate, fft_size, hop_size, return_complex=True) | |
| reference_mag = torch.abs(reference_stft) | |
| estimate_mag = torch.abs(estimated_stft) | |
| loss = 10 * F.l1_loss(estimate_mag, reference_mag) | |
| ret = 100 / (1. + float(loss.cpu().numpy())) | |
| return ret | |
| def LogWMSE_metric( | |
| reference: np.ndarray, | |
| estimate: np.ndarray, | |
| mixture: np.ndarray, | |
| device: str = 'cpu', | |
| ) -> float: | |
| """ | |
| Calculate the Log-WMSE (Logarithmic Weighted Mean Squared Error) between the reference, estimate, and mixture signals. | |
| This metric evaluates the quality of the estimated signal compared to the reference signal in the | |
| context of audio source separation. The result is given in logarithmic scale, which helps in evaluating | |
| signals with large amplitude differences. | |
| Parameters: | |
| ---------- | |
| reference : np.ndarray | |
| The ground truth audio signal of shape (channels, time), where channels is the number of audio channels | |
| (e.g., 1 for mono, 2 for stereo) and time is the length of the audio in samples. | |
| estimate : np.ndarray | |
| The estimated audio signal of shape (channels, time). | |
| mixture : np.ndarray | |
| The mixed audio signal of shape (channels, time). | |
| device : str, optional | |
| The device to run the computation on, either 'cpu' or 'cuda'. Default is 'cpu'. | |
| Returns: | |
| ------- | |
| float | |
| The Log-WMSE value, which quantifies the difference between the reference and estimated signal on a logarithmic scale. | |
| """ | |
| from torch_log_wmse import LogWMSE | |
| log_wmse = LogWMSE( | |
| audio_length=reference.shape[-1] / 44100, # audio length in seconds | |
| sample_rate=44100, # sample rate of 44100 Hz | |
| return_as_loss=False, # return as loss (False means return as metric) | |
| bypass_filter=False, # bypass frequency filtering (False means apply filter) | |
| ) | |
| reference = torch.from_numpy(reference).unsqueeze(0).unsqueeze(0).to(device) | |
| estimate = torch.from_numpy(estimate).unsqueeze(0).unsqueeze(0).to(device) | |
| mixture = torch.from_numpy(mixture).unsqueeze(0).to(device) | |
| res = log_wmse(mixture, reference, estimate) | |
| return float(res.cpu().numpy()) | |
| def AuraSTFT_metric( | |
| reference: np.ndarray, | |
| estimate: np.ndarray, | |
| device: str = 'cpu', | |
| ) -> float: | |
| """ | |
| Calculate the AuraSTFT metric, which evaluates the spectral difference between the reference and estimated | |
| audio signals using Short-Time Fourier Transform (STFT) loss. | |
| The AuraSTFT metric computes the STFT loss in both logarithmic and linear magnitudes, and it is commonly used | |
| to assess the quality of audio separation tasks. The result is returned as a value scaled to the range [0, 100]. | |
| Parameters: | |
| ---------- | |
| reference : np.ndarray | |
| The ground truth audio signal of shape (channels, time), where channels is the number of audio channels | |
| (e.g., 1 for mono, 2 for stereo) and time is the length of the audio in samples. | |
| estimate : np.ndarray | |
| The estimated audio signal of shape (channels, time). | |
| device : str, optional | |
| The device to run the computation on, either 'cpu' or 'cuda'. Default is 'cpu'. | |
| Returns: | |
| ------- | |
| float | |
| The AuraSTFT metric value, scaled to the range [0, 100], which quantifies the difference between | |
| the reference and estimated signal in the spectral domain. | |
| """ | |
| from auraloss.freq import STFTLoss | |
| stft_loss = STFTLoss( | |
| w_log_mag=1.0, # weight for log magnitude | |
| w_lin_mag=0.0, # weight for linear magnitude | |
| w_sc=1.0, # weight for spectral centroid | |
| device=device, | |
| ) | |
| reference = torch.from_numpy(reference).unsqueeze(0).to(device) | |
| estimate = torch.from_numpy(estimate).unsqueeze(0).to(device) | |
| res = 100 / (1. + 10 * stft_loss(reference, estimate)) | |
| return float(res.cpu().numpy()) | |
| def AuraMRSTFT_metric( | |
| reference: np.ndarray, | |
| estimate: np.ndarray, | |
| device: str = 'cpu', | |
| ) -> float: | |
| """ | |
| Calculate the AuraMRSTFT metric, which evaluates the spectral difference between the reference and estimated | |
| audio signals using Multi-Resolution Short-Time Fourier Transform (STFT) loss. | |
| The AuraMRSTFT metric uses multi-resolution STFT analysis, which allows better representation of both | |
| low- and high-frequency components in the audio signals. The result is returned as a value scaled to the range [0, 100]. | |
| Parameters: | |
| ---------- | |
| reference : np.ndarray | |
| The ground truth audio signal of shape (channels, time), where channels is the number of audio channels | |
| (e.g., 1 for mono, 2 for stereo) and time is the length of the audio in samples. | |
| estimate : np.ndarray | |
| The estimated audio signal of shape (channels, time). | |
| device : str, optional | |
| The device to run the computation on, either 'cpu' or 'cuda'. Default is 'cpu'. | |
| Returns: | |
| ------- | |
| float | |
| The AuraMRSTFT metric value, scaled to the range [0, 100], which quantifies the difference between | |
| the reference and estimated signal in the multi-resolution spectral domain. | |
| """ | |
| from auraloss.freq import MultiResolutionSTFTLoss | |
| mrstft_loss = MultiResolutionSTFTLoss( | |
| fft_sizes=[1024, 2048, 4096], | |
| hop_sizes=[256, 512, 1024], | |
| win_lengths=[1024, 2048, 4096], | |
| scale="mel", # mel scale for frequency resolution | |
| n_bins=128, # number of bins for mel scale | |
| sample_rate=44100, | |
| perceptual_weighting=True, # apply perceptual weighting | |
| device=device | |
| ) | |
| reference = torch.from_numpy(reference).unsqueeze(0).float().to(device) | |
| estimate = torch.from_numpy(estimate).unsqueeze(0).float().to(device) | |
| res = 100 / (1. + 10 * mrstft_loss(reference, estimate)) | |
| return float(res.cpu().numpy()) | |
| def bleed_full( | |
| reference: np.ndarray, | |
| estimate: np.ndarray, | |
| sr: int = 44100, | |
| n_fft: int = 4096, | |
| hop_length: int = 1024, | |
| n_mels: int = 512, | |
| device: str = 'cpu', | |
| ) -> Tuple[float, float]: | |
| """ | |
| Calculate the 'bleed' and 'fullness' metrics between a reference and an estimated audio signal. | |
| The 'bleed' metric measures how much the estimated signal bleeds into the reference signal, | |
| while the 'fullness' metric measures how much the estimated signal retains its distinctiveness | |
| in relation to the reference signal, both using mel spectrograms and decibel scaling. | |
| Parameters: | |
| ---------- | |
| reference : np.ndarray | |
| The reference audio signal, shape (channels, time), where channels is the number of audio channels | |
| (e.g., 1 for mono, 2 for stereo) and time is the length of the audio in samples. | |
| estimate : np.ndarray | |
| The estimated audio signal, shape (channels, time). | |
| sr : int, optional | |
| The sample rate of the audio signals. Default is 44100 Hz. | |
| n_fft : int, optional | |
| The FFT size used to compute the STFT. Default is 4096. | |
| hop_length : int, optional | |
| The hop length for STFT computation. Default is 1024. | |
| n_mels : int, optional | |
| The number of mel frequency bins. Default is 512. | |
| device : str, optional | |
| The device for computation, either 'cpu' or 'cuda'. Default is 'cpu'. | |
| Returns: | |
| ------- | |
| tuple | |
| A tuple containing two values: | |
| - `bleedless` (float): A score indicating how much 'bleeding' the estimated signal has (higher is better). | |
| - `fullness` (float): A score indicating how 'full' the estimated signal is (higher is better). | |
| """ | |
| from torchaudio.transforms import AmplitudeToDB | |
| reference = torch.from_numpy(reference).float().to(device) | |
| estimate = torch.from_numpy(estimate).float().to(device) | |
| window = torch.hann_window(n_fft).to(device) | |
| # Compute STFTs with the Hann window | |
| D1 = torch.abs(torch.stft(reference, n_fft=n_fft, hop_length=hop_length, window=window, return_complex=True, | |
| pad_mode="constant")) | |
| D2 = torch.abs(torch.stft(estimate, n_fft=n_fft, hop_length=hop_length, window=window, return_complex=True, | |
| pad_mode="constant")) | |
| mel_basis = librosa.filters.mel(sr=sr, n_fft=n_fft, n_mels=n_mels) | |
| mel_filter_bank = torch.from_numpy(mel_basis).to(device) | |
| S1_mel = torch.matmul(mel_filter_bank, D1) | |
| S2_mel = torch.matmul(mel_filter_bank, D2) | |
| S1_db = AmplitudeToDB(stype="magnitude", top_db=80)(S1_mel) | |
| S2_db = AmplitudeToDB(stype="magnitude", top_db=80)(S2_mel) | |
| diff = S2_db - S1_db | |
| positive_diff = diff[diff > 0] | |
| negative_diff = diff[diff < 0] | |
| average_positive = torch.mean(positive_diff) if positive_diff.numel() > 0 else torch.tensor(0.0).to(device) | |
| average_negative = torch.mean(negative_diff) if negative_diff.numel() > 0 else torch.tensor(0.0).to(device) | |
| bleedless = 100 * 1 / (average_positive + 1) | |
| fullness = 100 * 1 / (-average_negative + 1) | |
| return bleedless.cpu().numpy(), fullness.cpu().numpy() | |
| def get_metrics( | |
| metrics: List[str], | |
| reference: np.ndarray, | |
| estimate: np.ndarray, | |
| mix: np.ndarray, | |
| device: str = 'cpu', | |
| ) -> Dict[str, float]: | |
| """ | |
| Calculate a list of metrics to evaluate the performance of audio source separation models. | |
| The function computes the specified metrics based on the reference, estimate, and mixture. | |
| Parameters: | |
| ---------- | |
| metrics : List[str] | |
| A list of metric names to compute (e.g., ['sdr', 'si_sdr', 'l1_freq']). | |
| reference : np.ndarray | |
| The reference audio (true signal) with shape (channels, length). | |
| estimate : np.ndarray | |
| The estimated audio (predicted signal) with shape (channels, length). | |
| mix : np.ndarray | |
| The mixed audio signal with shape (channels, length). | |
| device : str, optional, default='cpu' | |
| The device ('cpu' or 'cuda') to perform the calculations on. | |
| Returns: | |
| ------- | |
| Dict[str, float] | |
| A dictionary containing the computed metric values. | |
| """ | |
| result = dict() | |
| # Adjust the length to be the same across all inputs | |
| min_length = min(reference.shape[1], estimate.shape[1]) | |
| reference = reference[..., :min_length] | |
| estimate = estimate[..., :min_length] | |
| mix = mix[..., :min_length] | |
| if 'sdr' in metrics: | |
| references = np.expand_dims(reference, axis=0) | |
| estimates = np.expand_dims(estimate, axis=0) | |
| result['sdr'] = sdr(references, estimates)[0] | |
| if 'si_sdr' in metrics: | |
| result['si_sdr'] = si_sdr(reference, estimate) | |
| if 'l1_freq' in metrics: | |
| result['l1_freq'] = L1Freq_metric(reference, estimate, device=device) | |
| if 'log_wmse' in metrics: | |
| result['log_wmse'] = LogWMSE_metric(reference, estimate, mix, device) | |
| if 'aura_stft' in metrics: | |
| result['aura_stft'] = AuraSTFT_metric(reference, estimate, device) | |
| if 'aura_mrstft' in metrics: | |
| result['aura_mrstft'] = AuraMRSTFT_metric(reference, estimate, device) | |
| if 'bleedless' in metrics or 'fullness' in metrics: | |
| bleedless, fullness = bleed_full(reference, estimate, device=device) | |
| if 'bleedless' in metrics: | |
| result['bleedless'] = bleedless | |
| if 'fullness' in metrics: | |
| result['fullness'] = fullness | |
| return result | |