generate_spectrogram.py

# python generate_spectrogram.py --audio-dir /path/to/audio_dir/ --spectrogram-dir /path/to/spectrogram_dir/

import argparse
from __future__ import division, print_function
from os import listdir
from os.path import isfile, join

from matplotlib import pyplot as plt
import scipy.io.wavfile as wav


def wav_to_spectrogram(audio_path, save_path, spectrogram_dimensions=(248, 248), noverlap=16, cmap='gray_r'):

    sample_rate, samples = wav.read(audio_path)

    fig = plt.figure()
    fig.set_size_inches((spectrogram_dimensions[0]/fig.get_dpi(), spectrogram_dimensions[1]/fig.get_dpi()))
    ax = plt.Axes(fig, [0., 0., 1., 1.])
    ax.set_axis_off()
    fig.add_axes(ax)
    ax.specgram(samples, cmap=cmap, Fs=2, noverlap=noverlap)
    ax.xaxis.set_major_locator(plt.NullLocator())
    ax.yaxis.set_major_locator(plt.NullLocator())
    fig.savefig(save_path, bbox_inches="tight", pad_inches=0)


def dir_to_spectrogram(audio_dir, spectrogram_dir, spectrogram_dimensions=(64, 64), noverlap=16, cmap='gray_r'):
    
    file_names = [f for f in listdir(audio_dir) if isfile(join(audio_dir, f)) and '.wav' in f]

    for file_name in file_names:
        print(file_name)
        audio_path = audio_dir + file_name
        spectogram_path = spectrogram_dir + file_name.replace('.wav', '.png')
        wav_to_spectrogram(audio_path, spectogram_path, spectrogram_dimensions=spectrogram_dimensions, noverlap=noverlap, cmap=cmap)

def get_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--audio-dir",
        default=None,
        type=str,
        help="if set, path where recordings are placed",
    )
    parser.add_argument(
        "--spectrogram-dir",
        default=None,
        type=str,
        help="if set, path where spectrogram to be placed",
    )

    return parser

if __name__ == '__main__':
    parser = get_parser()
    args = parser.parse_args()
    
    dir_to_spectrogram(audio_dir=args.audio_dir, spectrogram_dir=args.spectrogram_dir)