analyze.py

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
os.environ['CUDA_VISIBLE_DEVICES'] = ''

# try:
#     import tflite_runtime.interpreter as tflite
# except:
#     from tensorflow import lite as tflite
from tensorflow import lite as tflite

import argparse
import operator
import librosa
import numpy as np
import math
import time
import glob
import concurrent.futures
import copy
import sys, os

def loadModel():

    #global INPUT_LAYER_INDEX
    #global OUTPUT_LAYER_INDEX
    #global MDATA_INPUT_INDEX
    #global CLASSES

    print('LOADING TF LITE MODEL...', end=' ')

    mdlpath = os.path.join(os.path.dirname(os.path.abspath(sys.argv[0])), 'model', 'BirdNET_6K_GLOBAL_MODEL.tflite')
    # Load TFLite model and allocate tensors.
    interpreter = tflite.Interpreter(model_path=mdlpath)
    interpreter.allocate_tensors()

    # Get input and output tensors.
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()

    # Get input tensor index
    input_layer_index = input_details[0]['index']
    mdata_input_index = input_details[1]['index']
    output_layer_index = output_details[0]['index']

    # Load labels
    classes = []
    lblpath = os.path.join(os.path.dirname(os.path.abspath(sys.argv[0])), 'model', 'labels.txt')

    with open(lblpath, 'r') as lfile:
        for line in lfile.readlines():
            classes.append(line.replace('\n', ''))

    print('DONE!')

    return [input_layer_index, mdata_input_index, output_layer_index, classes, interpreter]

def loadCustomSpeciesList(path):

    slist = []
    if os.path.isfile(path):
        with open(path, 'r') as csfile:
            for line in csfile.readlines():
                slist.append(line.replace('\r', '').replace('\n', ''))
    else:
        raise Exception('Custom species list path does not exist!')

    return slist

def splitSignal(sig, rate, overlap, seconds=3.0, minlen=1.5):

    # Split signal with overlap
    sig_splits = []
    for i in range(0, len(sig), int((seconds - overlap) * rate)):
        split = sig[i:i + int(seconds * rate)]

        # End of signal?
        if len(split) < int(minlen * rate):
            break
        
        # Signal chunk too short? Fill with zeros.
        if len(split) < int(rate * seconds):
            temp = np.zeros((int(rate * seconds)))
            temp[:len(split)] = split
            split = temp
        
        sig_splits.append(split)

    return sig_splits

def readAudioData(path, overlap, sample_rate=48000):

    print('READING AUDIO DATA FROM FILE {}...'.format(os.path.split(path)[1]), end=' ', flush=True)

    # Open file with librosa (uses ffmpeg or libav)
    sig, rate = librosa.load(path, sr=sample_rate, mono=True, res_type='kaiser_fast')

    # Split audio into 3-second chunks
    chunks = splitSignal(sig, rate, overlap)

    print('DONE! READ {} CHUNKS.'.format(len(chunks)))

    return chunks

def convertMetadata(m):

    # Convert week to cosine
    if m[2] >= 1 and m[2] <= 48:
        m[2] = math.cos(math.radians(m[2] * 7.5)) + 1 
    else:
        m[2] = -1

    # Add binary mask
    mask = np.ones((3,))
    if m[0] == -1 or m[1] == -1:
        mask = np.zeros((3,))
    if m[2] == -1:
        mask[2] = 0.0

    return np.concatenate([m, mask])

def custom_sigmoid(x, sensitivity=1.0):
    return 1 / (1.0 + np.exp(-sensitivity * x))

def predict(sample, model, sensitivity):
    interpreter = model[4]
    # Make a prediction
    interpreter.set_tensor(model[0], np.array(sample[0], dtype='float32'))
    interpreter.set_tensor(model[1], np.array(sample[1], dtype='float32'))
    interpreter.invoke()
    prediction = interpreter.get_tensor(model[2])[0]

    # Apply custom sigmoid
    p_sigmoid = custom_sigmoid(prediction, sensitivity)

    # Get label and scores for pooled predictions
    p_labels = dict(zip(model[3], p_sigmoid))

    # Sort by score
    p_sorted = sorted(p_labels.items(), key=operator.itemgetter(1), reverse=True)

    # Remove species that are on blacklist
    for i in range(min(10, len(p_sorted))):
        if p_sorted[i][0] in ['Human_Human', 'Non-bird_Non-bird', 'Noise_Noise']:
            p_sorted[i] = (p_sorted[i][0], 0.0)

    # Only return first the top ten results
    return (p_sorted[:10], p_sigmoid)

def analyzeAudioData(chunks,args, sensitivity, model, file):

    # different format for standard and post-processing approach
    detections = {}
    detections_mea = np.zeros(shape=(len(chunks), len(model[3])))

    start = time.time()
    print('ANALYZING AUDIO FROM {} ...'.format(os.path.split(file)[1]), end=' ', flush=True)

    # Convert and prepare metadata
    mdata = convertMetadata(np.array([args.lat, args.lon, args.week]))
    mdata = np.expand_dims(mdata, 0)

    # Parse every chunk
    timestamps = []
    pred_start = 0.0

    i = 0
    for c in chunks:

        # Prepare as input signal
        sig = np.expand_dims(c, 0)

        # Make prediction
        p1, p2 = predict([sig, mdata], model, sensitivity)
        
        
        # Save result and timestamp
        detections_mea[i] = p2
        timestamps.append(pred_start)

        pred_end = pred_start + 3.0
        detections[file + ';' + str(pred_start) + ';' + str(pred_end)] = p1
        pred_start = pred_end - args.overlap
        i += 1

    print('DONE! TIME {:.1f} SECONDS'.format(time.time() - start))
    return (detections_mea.transpose(), timestamps, detections)

def writeDetectionsToFile(detections, path, white_list, file, args, sensitivity):

    print('WRITING RESULTS FROM {} TO {} ...'.format(os.path.split(file)[1], path), end=' ')
    rcnt = 0
    with open(path, 'a') as rfile:
        #rfile.write('Start (s);End (s);Scientific name;Common name;Confidence\n')
        for d in detections:
            for entry in detections[d]:
                if entry[1] >= args.min_conf and (entry[0] in white_list or len(white_list) == 0):
                    #rfile.write(d + ';' + entry[0].replace('_', ';') + ';' + str(entry[1]) + '\n')
                    rfile.write('{};{};{};{};{};{};{};{};{};{};FALSE\n'.format(d,entry[0].replace('_', ';'),entry[1], args.lat, args.lon, args.week, args.overlap, sensitivity, args.min_conf, args.custom_list if args.custom_list else 'NA'))
                    rcnt += 1
    print('DONE! WROTE {} RESULTS.'.format(rcnt))
    return rcnt

def concatFiles(filelist, output, output_mea, mea):
    """Concatenate different files from a list in one outputfile"""
    with open (output, 'a') as outfile:
        for fname in filelist:
            with open(fname[0]) as infile:
                for line in infile:
                    outfile.write(line)
    if mea:
        with open (output_mea, 'a') as outfile:
            for fname in filelist:
                with open(fname[1]) as infile:
                    for line in infile:
                        outfile.write(line)
                
    for files in filelist:
        # remove tempfiles
        os.remove(files[0])
        if mea:
            os.remove(files[1])

def initResultsFile(path):
    print("Initialize {} as result file".format(path))
    with open(path, "w") as file:
        file.write('Filepath;Start (s);End (s);Scientific name;Common name;Confidence;lat;lon;week;overlap;sensitivity;min_conf;custom_list;mea\n')
    return 0

def movingExpAverage(timetable, n=3):
    """calculate moving exponential average"""

    weights = np.exp(np.linspace(-1., 0., n))
    weights /= weights.sum()
    i = 0
    for row in timetable:
        # a = np.convolve(row, weights, mode='full')[:len(row)]
        a = np.convolve(row, weights, mode='full')[n-1:]
        # a[:n-1] = row[:n-1]
        # a[:n-1] = a[n-1]
        timetable[i] = a
        i += 1
    return timetable

def whiteListing(timetable, white_list, model):
    detections = {}
    i = 0
    for j in timetable:
        if (model[3][i] in white_list) or (len(white_list) == 0 and model[3][i] not in ['Human_Human', 'Non-bird_Non-bird', 'Noise_Noise']):
            detections[model[3][i]] = j
        i += 1
    return detections

def writeMeaToFile(detections, timestamps, path, file, args, sensitivity):

    print('WRITING RESULTS FROM {} TO {} ...'.format(os.path.split(file)[1], path), end=' ')
    with open(path, 'a') as rfile:  
        cntr = 0      
        for d in detections:
            rcnt = 0       
            for entry in detections[d]:
                if entry >= args.min_conf:
                    #rfile.write(file + ';' + str(timestamps[rcnt]) + ';' + str(timestamps[rcnt] + 3) + ';' + d.replace('_', ';') + ';' + str(entry) + '\n')
                    rfile.write('{};{};{};{};{};{};{};{};{};{};{};{};TRUE\n'.format(file, timestamps[rcnt], timestamps[rcnt] + 3, d.replace('_', ';'), entry, args.lat, args.lon, args.week, args.overlap, sensitivity, args.min_conf, args.custom_list if args.custom_list else 'NA'))
                    cntr += 1
                rcnt += 1
    print('DONE! WROTE {} RESULTS.'.format(cntr))
    return cntr

#def analyze(overlap, lat, lon, week, sensitivity, filelist, min_conf, i, white_list, mea):
def analyze(args, sensitivity, filelist, i, white_list):
    
    # generate tempfile-names 
    path="temp{}.csv".format(i)
    path_mea= "temp_mea{}.csv".format(i)

    # load Model
    model = loadModel()
    num_det = 0
    num_det_mea = 0

    for file in filelist:        
        # Read audio data
        audioData = readAudioData(file, args.overlap)

        # Process audio data and get detections
        pp_det, timestamps, def_det = analyzeAudioData(audioData, args, sensitivity, model, file)
        
        # Write standard detections (def_det) to tempfile
        num_det += writeDetectionsToFile(def_det, path, white_list, file, args, sensitivity)
        
        # apply moving exponential average to pp_det and write results to tempfile
        if args.mea:
            mea_det = whiteListing(movingExpAverage(pp_det), white_list, model)
            num_det_mea += writeMeaToFile(mea_det, timestamps, path_mea, file, args, sensitivity)

    return (path, path_mea, num_det, num_det_mea)

def main():
    a = time.time()
    # Parse passed arguments
    parser = argparse.ArgumentParser()
    parser.add_argument('--i', help='Path to input file.', required = True)
    parser.add_argument('--o', default='result.csv', help='Path to output file. Defaults to result.csv.')
    parser.add_argument('--lat', type=float, default=-1, help='Recording location latitude. Set -1 to ignore.')
    parser.add_argument('--lon', type=float, default=-1, help='Recording location longitude. Set -1 to ignore.')
    parser.add_argument('--week', type=int, default=-1, help='Week of the year when the recording was made. Values in [1, 48] (4 weeks per month). Set -1 to ignore.')
    parser.add_argument('--overlap', type=float, default=0.0, help='Overlap in seconds between extracted spectrograms. Values in [0.0, 2.9]. Defaults tp 0.0.')
    parser.add_argument('--sensitivity', type=float, default=1.0, help='Detection sensitivity; Higher values result in higher sensitivity. Values in [0.5, 1.5]. Defaults to 1.0.')
    parser.add_argument('--min_conf', type=float, default=0.1, help='Minimum confidence threshold. Values in [0.01, 0.99]. Defaults to 0.1.')   
    parser.add_argument('--custom_list', default='', help='Path to text file containing a list of species. Not used if not provided.')
    parser.add_argument('--threads', default=1, type=int, help='Integer; Number of threads')
    parser.add_argument('--recursive', default=False, choices = ['True', 'False'], help='Boolean; True, if the files should be searched recursively')
    parser.add_argument('--mea', default=False, choices = ['True', 'False'], help='Boolean; True, if moving exponential average of results should be calculated')

    args = parser.parse_args()
    try:
        # Load custom species list
        if not args.custom_list == '':
            white_list = loadCustomSpeciesList(args.custom_list)
        else:
            white_list = []

        #check variables range
        if not -90 <= args.lat <= 90:
            raise Exception('Argument --lat not in [-90, 90]')
        if not -180 <= args.overlap <= 180:
            raise Exception('Argument --lon not in [-180, 180]')
        if not 0 <= args.overlap <= 2.9:
            raise Exception('Argument --overlap not in [0.0, 2.9]')
        if not -1 <= args.week <= 48 or args.week == 0:
            raise Exception('Argument --week not in [1, 48]')
        if not 0.5 <= args.sensitivity <= 1.5:
            raise Exception('Argument --sensitivity not in [0.5, 1.5]')
        if not 0.01 <= args.min_conf <= 0.99:
            raise Exception('Argument --min_conf not in [0.01, 0.99]')

        # set variables
        # week = max(-1, min(args.week, 48))
        sensitivity = max(0.5, min(1.0 - (args.sensitivity - 1.0), 1.5))
        # min_conf = max(0.01, min(args.min_conf, 0.99))

        # create output file name for moving exponential average results
        mea_o = os.path.join(os.path.split(args.o)[0], 'mea_' + os.path.split(args.o)[1])

        
        if os.path.exists(args.o) or (args.mea and os.path.exists(mea_o)):
            raise Exception("Outputfile(s) {} and/or {} already exist! Please rename output file(s) and try again!".format(args.o, mea_o))           

        # create list of filenames 
        filelist = glob.glob(args.i, recursive=args.recursive)
        
        if not filelist:
            raise Exception("No file(s) to analyze. Check input file path and try again!")
        
        # create list of lists of filenames to pass to different threads
        step = -(-len(filelist)//int(args.threads))
        file_threads = [filelist[i:i + step] for i in range(0, len(filelist), step)]        
        
        # create result files
        initResultsFile(args.o)
        if args.mea:
            initResultsFile(mea_o)
        
        # run analyze on different threads    
        with concurrent.futures.ThreadPoolExecutor() as executer:
            #futures = [executer.submit(analyze, args.overlap, args.lat, args.lon, week, sensitivity, filelist, min_conf, i, copy.deepcopy(white_list), args.mea) for filelist, i in zip(file_threads, range(0,len(file_threads)))]
            futures = [executer.submit(analyze, args, sensitivity, filelist, i, copy.deepcopy(white_list)) for filelist, i in zip(file_threads, range(0,len(file_threads)))]
            results = [f.result() for f in futures]
            
            concatFiles(results, args.o, mea_o, args.mea)            
    
    except Exception as err:
        print('Error: {}'.format(err))
    
    else:
        e = time.time()
        print('#######################################')
        print('RESULTS:')
        print('TOTAL TIME: {:.1f} SECONDS'.format(e-a))
        print('ANALYZED FILES: {}'.format(len(filelist)))
        num_det = 0
        for i in results:
            num_det += i[2]
        print('{} DETECTIONS WRITTEN TO {}'.format(num_det, args.o))
        if args.mea:
            num_det_mea = 0
            for i in results:
                num_det_mea += i[3]
            print('{} DETECTIONS WRITTEN TO {}'.format(num_det_mea, mea_o))
        print('#######################################')
        print('USED ARGUMENTS (may be defaults):')
        print('--i: {}'.format(args.i))
        print('--o: {}'.format(args.o))
        print('--lat: {}'.format(args.lat))
        print('--lon: {}'.format(args.lon))
        print('--week: {}'.format(args.week))
        print('--overlap: {}'.format(args.overlap))
        print('--sensitivity: {}'.format(args.sensitivity))
        print('--min_conf: {}'.format(args.min_conf))
        print('--custom_list: {}'.format(args.custom_list))
        print('--threads: {}'.format(args.threads))
        print('--recursive: {}'.format(args.recursive))
        print('--mea: {}'.format(args.mea))        
        print('#######################################')
        
if __name__ == '__main__':

    main()

    # Example calls
    # python3 analyze.py --i '/absolute/path/to/your/soundfile/directory/*.wav' --lat 35.4244 --lon -120.7463 --week 18
    # python3 analyze.py --i 'relative/path/to/your/soundfile/directory/*.mp3' --lat 47.6766 --lon -122.294 --week 11 --overlap 1.5 --min_conf 0.25 --sensitivity 1.25 --custom_list 'example/custom_species_list.txt'