abtExp/source_separation

1# Make function to find classes in target directory2import os3import librosa4import torch5import numpy as np6from torchaudio.transforms import Resample7 8SAMPLE_RATE = 441009AUDIO_LEN = 2.9010 11# Parameters to control the MelSpec generation12N_MELS = 12813F_MIN = 2014F_MAX = 1600015N_FFT = 102416HOP_LEN = 51217 18# Make function to find classes in target directory19def find_classes(directory: str):20    # 1. Get the class names by scanning the target directory21    classes = sorted(entry.name for entry in os.scandir(directory) if entry.is_dir())22    # 2. Raise an error if class names not found23    if not classes:24        raise FileNotFoundError(f"Couldn't find any classes in {directory}.")25    # 3. Crearte a dictionary of index labels (computers prefer numerical rather than string labels)26    class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}27    return classes, class_to_idx28    29def resample(wav, sample_rate, new_sample_rate):30    if wav.shape[0] >= 2:31        wav = torch.mean(wav, dim=0)32    else:33        wav = wav.squeeze(0)34    if sample_rate > new_sample_rate:35        resampler = Resample(sample_rate, new_sample_rate)36        wav = resampler(wav)37    return wav38 39def mono_to_color(X, eps=1e-6, mean=None, std=None):40    X = np.stack([X, X, X], axis=-1)41    # Standardize42    mean = mean or X.mean()43    std = std or X.std()44    X = (X - mean) / (std + eps)45    # Normalize to [0, 255]46    _min, _max = X.min(), X.max()47    if (_max - _min) > eps:48        V = np.clip(X, _min, _max)49        V = 255 * (V - _min) / (_max - _min)50        V = V.astype(np.uint8)51    else:52        V = np.zeros_like(X, dtype=np.uint8)53    return V54 55def normalize(image, mean=None, std=None):56    image = image / 255.057    if mean is not None and std is not None:58        image = (image - mean) / std59    return np.moveaxis(image, 2, 0).astype(np.float32)60 61def compute_melspec(wav, sample_rate=SAMPLE_RATE):62    melspec = librosa.feature.melspectrogram(63        y=wav,64        sr=sample_rate, 65        n_fft=N_FFT, 66        fmin=F_MIN, 67        fmax=F_MAX,68        n_mels=N_MELS,69        hop_length=HOP_LEN70    )71    melspec = librosa.power_to_db(melspec).astype(np.float32)72    return melspec73 74def audio_preprocess(wav, sample_rate):75    wav = wav.numpy()76    melspec = compute_melspec(wav, sample_rate)77    image = mono_to_color(melspec)78    image = normalize(image, mean=None, std=None)79    image = torch.from_numpy(image)80    return image