src/preprocess.py

"""
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)