compare_codec/spectrogram.py

"""Generate mel spectrogram images from audio arrays."""

from __future__ import annotations

import tempfile
from pathlib import Path

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


def generate(audio: np.ndarray, sr: int) -> Path:
    """Render a mel spectrogram to a temporary PNG and return its path."""
    fig, ax = plt.subplots(1, 1, figsize=(8, 3))
    S = librosa.feature.melspectrogram(y=audio, sr=sr, n_mels=128)
    S_dB = librosa.power_to_db(S, ref=np.max)
    librosa.display.specshow(S_dB, sr=sr, x_axis="time", y_axis="mel", ax=ax)
    ax.set(title=None, xlabel=None, ylabel=None)
    ax.tick_params(labelsize=8)
    fig.tight_layout(pad=0.5)

    path = Path(tempfile.mktemp(suffix=".png"))
    fig.savefig(path, dpi=100, bbox_inches="tight")
    plt.close(fig)
    return path