""" Copyright (c) : Calixte Mayoraz 2024 https://gitlab.com/calixtemayoraz Preprocessing for use in the Music-Squid model. Regarding Chunk sizes: Since the mel spectrogram uses 512-long windows for its fft, if we want a 128-long mel spectrogram, we need 128 x 512 = 65536 samples. At 16000 samples per second, this corresponds to 4.096s. We want to use 8 slices (to get some nice powers of 2 in there) so we need 8 x 4.096 = 32.768 or 524288 samples. """ import logging import librosa import numpy as np import os SAMPLE_RATE = 16000 FULL_CHUNK_TIME = 32.768 # full chunk, not single slice! SINGLE_SLICE_TIME = 4.096 SAMPLES_PER_CHUNK = int(SAMPLE_RATE * FULL_CHUNK_TIME) SAMPLES_PER_SLICE = int(SAMPLE_RATE * SINGLE_SLICE_TIME) CACHE_DIR = os.path.abspath(os.path.join(__file__, '..', '..', '.cache')) if not os.path.exists(CACHE_DIR): os.mkdir(CACHE_DIR) logging.info("using cache dir %s", CACHE_DIR) def load(filepath: str) -> np.ndarray: """ Load filepath into mono with a sample rate of 16000. saves the output to cache before returning it Parameters ---------- filepath: str the file to load Returns ------- np.ndarray 1D array of the loaded audio file """ cache_path = os.path.join(CACHE_DIR, os.path.split(filepath)[-1]) cache_path = os.path.splitext(cache_path)[0] + ".npy" if os.path.exists(cache_path): return np.load(cache_path) output = librosa.load(filepath, sr=SAMPLE_RATE)[0] np.save(cache_path, output) return output def signal_to_mel(data) -> np.ndarray: """ Prepares the Mel Spectrogram for the model Parameters ---------- data: np.ndarray Returns ------- np.ndarray """ # get a mel spectrogram for everything mel = librosa.feature.melspectrogram(y=data, sr=SAMPLE_RATE) # slice into 128x128 chunks mel = np.squeeze(np.lib.stride_tricks.sliding_window_view(mel, (128, 128)))[::128] # now normalize each chunk mel = np.log10(mel + 1e-10) mel = (mel - np.min(mel, axis=(1, 2))[:, np.newaxis, np.newaxis]) / ( np.max(mel, axis=(1, 2))[:, np.newaxis, np.newaxis] - np.min(mel, axis=(1, 2))[:, np.newaxis, np.newaxis]) return np.nan_to_num(np.expand_dims(mel, -1), nan=0) def signal_to_tempogram(data) -> np.ndarray: """ Prepares the tempogram for the model Parameters ---------- data: np.ndarray Returns ------- np.ndarray """ # tempogram the whole thing (keep only the first 128 tempo rows) tempo = librosa.feature.tempogram(y=data, sr=SAMPLE_RATE)[:128] # slice into 4.096s slices tempo = np.squeeze(np.lib.stride_tricks.sliding_window_view(tempo, (128, 128)))[::128] # average each slice column wise tempo = (tempo - np.min(tempo, axis=1)[:, np.newaxis]) / ( np.max(tempo, axis=1)[:, np.newaxis] - np.min(tempo, axis=1)[:, np.newaxis]) tempo = np.nan_to_num(tempo, nan=0) tempo = np.expand_dims(np.squeeze(np.lib.stride_tricks.sliding_window_view(tempo, (8, 128, 128))), -1) tempo = np.mean(tempo, axis=-2) return (tempo - np.mean(tempo, axis=(1,2,3), keepdims=True)) / np.std(tempo, axis=(1,2,3), keepdims=True) def signal_to_chromagram(data): """ Prepares the chromagram for the model Parameters ---------- data: np.ndarray Returns ------- np.ndarray """ # chromagram the whole thing chroma = librosa.feature.chroma_cqt(y=data, sr=SAMPLE_RATE) # slice into slices chroma = np.squeeze(np.lib.stride_tricks.sliding_window_view(chroma, (12, 128)))[::128] # normalize each slice chroma = (chroma - np.min(chroma, axis=1)[:, np.newaxis]) / ( np.max(chroma, axis=1)[:, np.newaxis] - np.min(chroma, axis=1)[:, np.newaxis]) chroma = np.nan_to_num(chroma, nan=0) chroma = np.expand_dims(np.squeeze(np.lib.stride_tricks.sliding_window_view(chroma, (8, 12, 128))), -1) chroma = np.mean(chroma, axis=-2) return (chroma - np.mean(chroma, axis=(1, 2, 3), keepdims=True)) / np.std(chroma, axis=(1, 2, 3), keepdims=True) def preprocess(filename) -> 'tuple[np.ndarray, np.ndarray, np.ndarray]': """ Preprocesses the file to input ready to be plugged into the model. Since the model takes in 30s pieces of a track, the preprocessing returns the following: mel spectrogram: (n, 8, 128, 128, 1) float16 tempogram: (n, 8, 128, 1) float16 chromagram: (n, 8, 12, 1) float16 Parameters ---------- filename: str Returns ------- tuple[np.ndarray, np.ndarray, np.ndarray] """ data = load(filename) mel, tempo, chroma = signal_to_mel(data), signal_to_tempogram(data), signal_to_chromagram(data) mel = np.expand_dims(np.squeeze(np.lib.stride_tricks.sliding_window_view(mel, (8, 128, 128, 1))),-1) return mel.astype(np.float16), tempo.astype(np.float16), chroma.astype(np.float16)