display.py

import matplotlib.pyplot as plt
import librosa.display


def _display_stft(stft, label, *, sample_rate=16000, hop_length=128, ax=None, y_axis='linear'):
    if ax is None:
        raise ValueError("ax must be provided.")

    librosa.display.specshow(stft, sr=sample_rate, hop_length=hop_length, x_axis='time', y_axis=y_axis, ax=ax)
    ax.set_ylabel("Frequency (Hz)")
    ax.set_xlabel("Time (s)")
    ax.set_title(f"STFT: {label}")


def display_stft(tf_record, n=4, *, sample_rate=16000, hop_length=128, y_axis='linear'):
    if n > 3:
        raise ValueError("n must be less than 3.")
    rows = n
    cols = n
    n = n * n
    fig, ax = plt.subplots(rows, cols, figsize=(18, rows * 6))
    i, j = 0, 0
    for stft, label in tf_record.take(n).as_numpy_iterator():
        if j >= cols:
            i += 1
            j = 0

        if cols == 1 and rows == 1:
            _display_stft(stft, label, sample_rate=sample_rate,
                          hop_length=hop_length, ax=ax, y_axis=y_axis)
        if cols == 1:
            _display_stft(stft, label, sample_rate=sample_rate,
                          hop_length=hop_length, ax=ax[i], y_axis=y_axis)
        if rows == 1:
            _display_stft(stft, label, sample_rate=sample_rate,
                          hop_length=hop_length, ax=ax[j], y_axis=y_axis)
        else:
            _display_stft(stft, label, sample_rate=sample_rate, hop_length=hop_length, y_axis=y_axis,
                          ax=ax[i, j])

        j += 1

    while j < cols:
        ax[i, j].axis('off')
        j += 1

    fig.subplots_adjust(bottom=0.3, top=0.8, wspace=0.3, hspace=0.5)

    plt.show()