Skip to content
Snippets Groups Projects
Select Git revision
  • 7dfceb6084e501b82f38a2cfbd9a4d3352cb3874
  • main default protected
  • V1
3 results

utils.py

Blame
  • 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')