Select Git revision
get_spectrogram.py 4.28 KiB
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')