inference_ensemble.py

import argparse
from tools.AFibInference import AFibInference
import numpy as np
import os
from math import floor
import pandas as pd
import pywt

def parse_command_line_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_path', type=str, required=False, help='The path to the input document in '
                                                                       'csv format')
    parser.add_argument('--model_name', type=str, required=False, help='The name of the model')
    parser.add_argument('--data_path', type=str, required=False, help='The path to the input document')
    parser.add_argument('--threshold', type=float, default=0.5, help='The threshold for the prediction')
    parser.add_argument('--use_torch', type=str, default=True, help='Whether to use PyTorch')
    return parser.parse_args()

def inference_torch(classifier, data, threshold):
    import torch
    data = torch.tensor(data, dtype=torch.float32)
    return classifier.classify(data, threshold=threshold)

def inference_onnx(classifier, data, threshold):
    data = np.array(data, dtype=np.float32)
    predictions = classifier.classify(data=data, threshold=threshold)
    return predictions

def filtering(data):
    wavelet_type = 'db10'
    data_wt = np.zeros((data.shape[0], 1, 1280))
    for i in range(data.shape[0]):
        signal = data[i,0,:]
        DWTcoeffs = pywt.wavedec(signal, wavelet_type, mode='symmetric', level=4, axis=-1)
        signal_filt = pywt.waverec(DWTcoeffs, wavelet_type, mode='symmetric', axis=-1)
        data_wt[i, 0, :] = signal
    return data_wt

def split_signal(data, length=10., sampling_freq=128.):
    div = floor(data.shape[0] / (sampling_freq*length))
    data = data[:int(div*length*sampling_freq)]
    data = [[i] for i in np.split(data, div)]
    return np.array(data)

def from_csv_to_numpy(file_name):
    data_file = pd.read_csv(file_name)
    data_file = data_file['value'].to_numpy()
    return data_file
import torch
def main():
    args = parse_command_line_arguments()
    use_torch = True if args.use_torch == 'True' else False
    model_path = args.model_path
    model_name = args.model_name
    data_path = args.data_path
    threshold = args.threshold
    data = from_csv_to_numpy(data_path)
    signals = filtering(split_signal(data))
    classifier = AFibInference(model_name=model_name, weights_path=model_path, deep_ensemble=True, use_torch=use_torch)
    if use_torch:
        classification, uncertainty = inference_torch(classifier=classifier, threshold=threshold, data=signals)
    else:
        classification, uncertainty = inference_onnx(classifier=classifier, threshold=threshold, data=signals)
    print(classification)
    print(uncertainty)


if __name__ == '__main__':
    main()