Spaces:
Running
Running
| # This file contains code adapted from the following sources: | |
| # [MIT license] https://github.com/mir-aidj/all-in-one/blob/main/src/allin1/postprocessing/helpers.py | |
| import numpy as np | |
| import torch.nn.functional as F | |
| import torch | |
| import librosa | |
| from typing import Union | |
| from scipy.signal import argrelextrema | |
| from scipy.interpolate import interp1d | |
| from numpy.lib.stride_tricks import sliding_window_view | |
| from numpy.typing import NDArray | |
| def local_maxima(tensor, filter_size=41): | |
| assert len(tensor.shape) in (1, 2), "Input tensor should have 1 or 2 dimensions" | |
| assert filter_size % 2 == 1, "Filter size should be an odd number" | |
| original_shape = tensor.shape | |
| if len(original_shape) == 1: | |
| tensor = tensor.unsqueeze(0) | |
| # Pad the input array with the minimum value | |
| padding = filter_size // 2 | |
| padded_arr = F.pad(tensor, (padding, padding), mode="constant", value=-torch.inf) | |
| # Create a rolling window view of the padded array | |
| rolling_view = padded_arr.unfold(1, filter_size, 1) | |
| # Find the indices of the local maxima | |
| center = filter_size // 2 | |
| local_maxima_mask = torch.eq( | |
| rolling_view[:, :, center], torch.max(rolling_view, dim=-1).values | |
| ) | |
| local_maxima_indices = local_maxima_mask.nonzero() | |
| # Initialize a new PyTorch tensor with zeros and the same shape as the input tensor | |
| output_arr = torch.zeros_like(tensor) | |
| # Set the local maxima values in the output tensor | |
| output_arr[local_maxima_mask] = tensor[local_maxima_mask] | |
| output_arr = output_arr.reshape(original_shape) | |
| return output_arr, local_maxima_indices | |
| def local_maxima_numpy(arr, order=20): | |
| is_batch = len(arr.shape) == 2 | |
| if is_batch: | |
| return np.stack([local_maxima_numpy(x, order) for x in arr]) | |
| # Define a comparison function for argrelextrema to find local maxima | |
| compare_func = np.greater | |
| # Find the indices of the local maxima | |
| local_maxima_indices = argrelextrema(arr, compare_func, order=order) | |
| # Initialize a new numpy array with zeros and the same shape as the input array | |
| output_arr = np.zeros_like(arr) | |
| # Set the local maxima values in the output array | |
| output_arr[local_maxima_indices] = arr[local_maxima_indices] | |
| return output_arr | |
| def peak_picking(boundary_activation, window_past=12, window_future=6): | |
| # Find local maxima using a sliding window | |
| window_size = window_past + window_future | |
| assert window_size % 2 == 0, "window_past + window_future must be even" | |
| window_size += 1 | |
| # Pad boundary_activation | |
| boundary_activation_padded = np.pad( | |
| boundary_activation, (window_past, window_future), mode="constant" | |
| ) | |
| max_filter = sliding_window_view(boundary_activation_padded, window_size) | |
| local_maxima = (boundary_activation == np.max(max_filter, axis=-1)) & ( | |
| boundary_activation > 0 | |
| ) | |
| # Compute strength values by subtracting the mean of the past and future windows | |
| past_window_filter = sliding_window_view( | |
| boundary_activation_padded[: -(window_future + 1)], window_past | |
| ) | |
| future_window_filter = sliding_window_view( | |
| boundary_activation_padded[window_past + 1 :], window_future | |
| ) | |
| past_mean = np.mean(past_window_filter, axis=-1) | |
| future_mean = np.mean(future_window_filter, axis=-1) | |
| strength_values = boundary_activation - ((past_mean + future_mean) / 2) | |
| # Get boundary candidates and their corresponding strength values | |
| boundary_candidates = np.flatnonzero(local_maxima) | |
| strength_values = strength_values[boundary_candidates] | |
| strength_activations = np.zeros_like(boundary_activation) | |
| strength_activations[boundary_candidates] = strength_values | |
| return strength_activations | |