utils.py

import numpy as np
import librosa
import matplotlib.pyplot as plt
from librosa.display import specshow


def plot_spec(y: np.ndarray, sr: int, title: str = "Spectrogram") -> plt.Figure:
    y[np.isnan(y)] = 0
    y[np.isinf(y)] = 0
    stft = librosa.stft(y=y)
    D = librosa.amplitude_to_db(np.abs(stft), ref=np.max)

    fig = plt.figure(figsize=(10, 4))
    specshow(D, sr=sr, y_axis="linear", x_axis="time", cmap="viridis")
    plt.title(title)
    plt.tight_layout()

    return fig