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

get_spectrogram.py

Blame
  • stephane.chavin's avatar
    Stephane Chavin authored
    9ca6ea46
    History
    get_spectrogram.py 2.78 KiB
    import os
    import pandas as pd
    import librosa
    import numpy as np
    import matplotlib.pyplot as plt
    from tqdm import tqdm
    import ipdb
    import argparse
    from p_tqdm import p_map
    import warnings
    warnings.filterwarnings('ignore')
    
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter, description='TODO')
    parser.add_argument('-f','--file', type=str,required=True,help = 'Name of the file that contain the recording to print')
    parser.add_argument('-p','--path_to_data', type=str, help = 'Path of the folder that contain the recodings',required=True)
    parser.add_argument('-d','--direction', type=str, help = 'Directory to wich spectrogram will be stored',required=True)
    parser.add_argument('-m','--mode', type=str,choose={'unique','multiple'}, help = 'Direction of the saved spectrogram',required=True)
    args = parser.parse_args()
    
    path_to_data = args.path_to_data
    direction = args.direction
    folder = 'Spectrogram/'
    DURATION = 5
    OVERLAP = 2
    NB_IMG_PER_REC = 1
    
    count = 0
    
    df = pd.read_csv(str(args.file),low_memory=False)
    df['Path'] = ''
    for w in range (len(df)):
        df.loc[w,'Path']=str(args.path_to_data+ df.code_unique[w]+"_split_" + str(df.chunk_initial_time[w]) + "_" + str(df.chunk_final_time[w]))
    
    
    def process(x):
       count, (i) = x
    
       for j in range (NB_IMG_PER_REC): #30*8 secondes - 30*2 secondes (overlap) = 180 secondes affichées sur 30 images : n'affiche que les 3 premières minutes d'un enregistrement
    
            filename = str(i[0])
    
            if args.mode == 'multiple' :
                if count ==0:
                    duration = DURATION
                    offset = 0
                else :
                    duration = DURATION
                    offset = offset + (DURATION - OVERLAP)
            if args.mode == 'unique':
                offset = 0
    
            try:
                y, sr = librosa.load(filename, offset = offset, duration = DURATION)
            except Exception:
                continue
            window_size = 1024
            window = np.hanning(window_size)
            stft  = librosa.core.spectrum.stft(y, n_fft=window_size, hop_length=512, window=window)
            out = 2 * np.abs(stft) / np.sum(window)
    
            plt.close()
            plt.figure()
    
            vmin = np.flipud(np.log10(np.abs(stft))).mean()
            vmax = np.flipud(np.log10(np.abs(stft))).max()
    
            plt.imshow(np.flipud(np.log10(np.abs(stft))),aspect = "auto",interpolation = None,cmap = 'jet',vmin = vmin,vmax = vmax)
    
            plt.subplots_adjust(top=1, bottom=0, left=0, right=1)
    
            name = str(i[0].replace('/','_').split('.')[0]+'_'+str(count))
            try :
    
                plt.savefig(str(direction+folder+name+'.jpg'))
    
            except FileNotFoundError:
    
                os.mkdir(str(direction+folder))
                plt.savefig(str(direction+folder+name+'.jpg'))
    
    p_map(process, enumerate(df.groupby('Path')), num_cpus=10) #la colonne dir correspond au path