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

get_time_freq_detection.py

Blame
  • Stephane Chavin's avatar
    Stephane Chavin authored
    2afd11a9
    History
    get_time_freq_detection.py 2.61 KiB
    import pandas as pd
    import os
    import argparse
    from datetime import date
    from tqdm import tqdm
    
    def arg_directory(path):
        if os.path.isdir(path):
            return path
        else:
            raise argparse.ArgumentTypeError(f'`{path}` is not a valid path')
    
    def process_annotations(annotations_folder, duration, sr):
        today = date.today()
        out_file = f'YOLO_detection_{today.day}_{today.month}'
    
        df_list = []
        names = []  # Add your class names here
    
        for file_name in tqdm(os.listdir(annotations_folder)):
            if file_name.endswith('.txt'):
                file_path = os.path.join(annotations_folder, file_name)
                annotation_df = pd.read_csv(file_path, sep=' ', names=['espece', 'x', 'y', 'w', 'h'])
    
                annotation_df['file'] = file_name
                annotation_df['idx'] = annotation_df['file'].str.split('_').str[-1].str.split('.').str[0]
                annotation_df['file'] = annotation_df['file'].str.rsplit('.', 1).str[0] + '.wav'
    
                annotation_df['annot'] = annotation_df['espece'].apply(lambda x: names[x])
    
                annotation_df['midl'] = (annotation_df['x'] * duration) + annotation_df['idx'].astype(int)
                annotation_df['freq_center'] = (1 - annotation_df['y']) * (sr / 2)
                annotation_df['freq_min'] = annotation_df['freq_center'] - (annotation_df['h'] * (sr / 2)) / 2
                annotation_df['freq_max'] = annotation_df['freq_center'] + (annotation_df['h'] * (sr / 2)) / 2
                annotation_df['start'] = annotation_df['midl'] - (annotation_df['w'] * duration) / 2
                annotation_df['stop'] = annotation_df['midl'] + (annotation_df['w'] * duration) / 2
                annotation_df['duration'] = annotation_df['stop'] - annotation_df['start']
    
                df_list.append(annotation_df)
    
        result_df = pd.concat(df_list, ignore_index=True)
        result_df.to_csv(os.path.join(outdir, f'{out_file}.csv'), index=False)
        print(f'Saved as {os.path.join(outdir, f"{out_file}.csv")}')
    
    if __name__ == "__main__":
        parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter, description='TODO')
        parser.add_argument('-p', '--path_to_data', type=arg_directory, help='Path of the folder that contains the .txt files', required=True)
        parser.add_argument('-d', '--directory', type=arg_directory, help='Directory where the dataframe will be stored', required=True)
        parser.add_argument('-t', '--duration', type=int, help='Duration of the spectrogram', required=True)
        parser.add_argument('-s', '--SR', type=int, help='Sampling Rate of the spectrogram')
        args = parser.parse_args()
    
        process_annotations(args.path_to_data, args.duration, args.SR)