ecg_feature_extraction.py

"""This file contains the ecg feature extraction pipelines and functions used for calculating the features."""

import neurokit2 as nk
from .configs import selected_features


def get_hrv_features(ecg_signal, fs):
    # Find peaks
    peaks, info = nk.ecg_peaks(ecg_signal, sampling_rate=fs, method="pantompkins1985")

    # Compute time domain features
    hrv_time_features = nk.hrv_time(peaks, sampling_rate=fs)

    # Compute frequency domain features
    #hrv_frequency_features = nk.hrv_frequency(peaks, sampling_rate=fs, method="welch", show=False)

    # Concat features
    #hrv_features = pd.concat([hrv_time_features, hrv_frequency_features], axis=1)
    hrv_features = hrv_time_features

    # to dict
    hrv_features = hrv_features[selected_features].to_dict(orient="records")

    return hrv_features


def normalize_features(features_df, normalization_method):
    if normalization_method == "difference":
        baseline_features = features_df[features_df['baseline'] == True].iloc[0]
        features_df.loc[features_df['baseline'] == False, features_df.columns.isin(selected_features)] -= baseline_features
    elif normalization_method == "relative":
        baseline_features = features_df[features_df['baseline'] == True].iloc[0]
        features_df.loc[features_df['baseline'] == False, features_df.columns.isin(selected_features)] /= baseline_features
    elif (normalization_method == "separate") or (normalization_method is None):
        pass

    return features_df