get_spectrogram.py

import os
import librosa
import glob
import argparse
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from p_tqdm import p_map

import warnings
warnings.filterwarnings('ignore')

def arg_directory(path):
    if os.path.isdir(path):
        return path
    else:
        raise argparse.ArgumentTypeError(f'`{path}` is not a valid path')

def create_spectrogram(y, directory, filename, offset, duration, window_arg, hop_length_arg):
    window = np.hanning(window_arg)
    stft = librosa.core.spectrum.stft(y, n_fft=window_arg, hop_length=hop_length_arg, window=window)

    plt.close()
    plt.figure()

    log_stft = np.log10(np.abs(stft))
    vmin, vmax = log_stft.min(), log_stft.max()

    plt.imshow(log_stft[::-1], aspect="auto", interpolation=None, cmap='jet', vmin=vmin, vmax=vmax)
    plt.subplots_adjust(top=1, bottom=0, left=0, right=1)

    name = os.path.join(directory, 'Spectrogram', f"{filename.replace('/', '_').split('.')[0]}_{offset}")
    
    try:
        plt.savefig(name + '.jpg')
    except FileNotFoundError:
        os.makedirs(os.path.join(directory, 'Spectrogram'), exist_ok=True)
        plt.savefig(name + '.jpg')

def process_recordings(data, img_per_rec, args):
    _, (i) = data
    duration = args.duration
    overlap = args.overlap

    for count in range(args.img_per_rec):
        offset = count * (duration - overlap)
        filename = str(i[0])
        
        try:
            y, _ = librosa.load(filename, offset=offset, duration=duration, sr=args.sr)
            create_spectrogram(y, args.directory, filename, offset, duration, args.window, args.hop)
        except Exception:
            print(f'`{filename}` cannot be open...')

if __name__ == "__main__":
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter, description='Extract spectrogram for each .wav file')
    parser.add_argument('path_to_data', type=arg_directory, help='Path of the folder that contains the recordings', required=True)
    parser.add_argument('directory', type=arg_directory, help='Directory to which spectrograms will be stored', required=True)
    parser.add_argument('-m', '--mode', type=str, choices=['unique', 'multiple'], help='Direction of the saved spectrogram', default='multiple')
    parser.add_argument('-n', '--columns_name', type=str, help='Name of the columns that contain the path of the .wav', default='Path')
    parser.add_argument('-i', '--input', type=str, choices=['file', 'folder'], help='Choose "file" if you have a .csv file or "folder" to export '
                                                                            'spectrogram from all the .wav of a folder', default='folder')
    parser.add_argument('-f', '--file', type=str, help='Name of the file that contains the recording to print', default=None)
    parser.add_argument('--frames', type=int, help='Number of spectrogram per file', default=30)
    parser.add_argument('--duration', type=int, help='Duration for each spectrogram', default=8)
    parser.add_argument('--overlap', type=int, help='Overlap between 2 spectrograms', default=2)
    parser.add_argument('--sr', type=int, help='Sampling rate for the spectrogram. If no argument, '
                                            'SR will be original SR of the recording', default=None)
    parser.add_argument('--window', type=int, help='Window size for the Fourier Transform', default=1024)
    parser.add_argument('--hop', type=int, help='Hop lenght for the Fourier Transform', default=512)
    parser.add_argument('--cpu', type=int, help='To speed up the process, write 2 or more', default=1)
    args = parser.parse_args()

    if args.mode == 'multiple':
        img_per_rec = args.frames
    elif args.mode == 'unique':
        img_per_rec = 1

    path_to_data = args.path_to_data

    if args.input == 'file':
        df = pd.read_csv(args.file, low_memory=False)
        df['Path'] = df[args.columns_name]
    elif args.input == 'folder':
        df = pd.DataFrame(glob.glob(os.path.join(path_to_data, '*'), recursive=True), columns=['Path'])

    img_per_rec = [img_per_rec]*len(df.groupby('Path'))
    args = [args]*len(df.groupby('Path'))    
    
    p_map(process_recordings, enumerate(df.groupby('Path')), img_per_rec, args, num_cpus=args[0].cpu, total=len(df.groupby('Path')))
    print(f'Saved to {args[0].directory}/Spectrogram')