get_time_freq_detection.py

"""Compiles detections into a dataframe and/or into Raven annotation format (.txt)"""

import argparse
import os
import yaml
import utils
import xarray as xr
import pandas as pd


def main(arguments):
    """
    Load all the informations to concatenate detection and get time/frequency informations
    and save them into a full file and multiple raven annotation file if --raven
    :param arguments (args): All the arguments of the script
    """
    with open(arguments.names, 'r', encoding='utf-8') as file:
        data = yaml.safe_load(file)
    names = data['names']
    df, dir_path = utils.detection2time_freq(annotations_folder=arguments.path_to_data,
                                             duration=arguments.duration,
                                             outdir=arguments.directory,
                                             sr=arguments.sr,
                                             names=names,
                                             wav=args.path_to_wav,
                                             raven=args.raven)
    # Convert DataFrame to xarray Dataset
    ds = xr.Dataset.from_dataframe(df)

    # Add metadata attributes
    ds.attrs['title'] = 'YOLO detection final table'
    ds.attrs['summary'] = 'This dataset contains the YOLO detection with parameters.'
    ds.attrs['description'] = str('The data includes positions : "pos", frequencies :'
                             ' "low freq (hz)"; "high freq (hz)", and timings : '
                             '"begin time (s)"; "end time (s)" and the species detected'
                             ' in the given files.')
    ds.attrs['date_created'] = pd.Timestamp.now().isoformat()
    ds.attrs['creator_name'] = os.getlogin()

    # Save Dataset to NetCDF file
    ds.to_netcdf(dir_path)
    print(f'Saved as {dir_path}')


if __name__ == "__main__":
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                     description='Collect detections and return a complete dataframe')
    parser.add_argument('path_to_data', type=utils.arg_directory,
                        help='Path of the folder that contains the .txt files')
    parser.add_argument('directory', type=utils.arg_directory,
                        help='Directory where the dataframe will be stored')
    parser.add_argument('names', type=str,
                        help='path to YOLOv5 custom_data.yaml file')
    parser.add_argument('-s', '--sr', type=int,
                        help='Sampling Rate of the spectrogram', required=True)
    parser.add_argument('--duration', type=int,
                        help='Duration of the spectrogram', default=8)
    parser.add_argument('--path_to_wav', type=utils.arg_directory,
                        help='Path of the folder that contains the .wav files', required=True)
    parser.add_argument('--raven', action='store_const', const=1, default=None,
                        help='Export the .txt per .WAV file to vizualize on Raven')
    args = parser.parse_args()

    main(args)