Authors: Danial Beg (dbeg@uci.edu) and Sahithi Chimmula (schimmul@uci.edu)
This project aims to detect atrial fibrilation (AFib) utilizing ECG waveform data provided in the PTB-XL dataset. A write-up of this project can be found here and please feel free to check out our final presentation
Some of the datasets are too big to include on GitHub and thus can be downloaded from the following links: ecgeq-500hzsrfava.npy and af_dataset.csv.
Once downloaded, please put them into the data folder to run the notebooks.
Additionally, please download output_file.pt from Google Drive and place it under src/transformer_data/.
AFib-Detection-Survey/
│
├── src/ # Source files
│ ├── Inference Notebok.ipynb # Main script running inference through saved model files
│ └── LSTM_Training.ipynb # How the LSTM training was done
│ └── MLP_and_LogRegression.ipynb # How the MLP and Logistic Regression training was done
│ └── Final_Transformer.ipynb # How the Transformer training was done (unweighted)
│ └── Final_Transformer_weighted.ipynb # How the Transformer training was done (weighted)
│ └── Pre_processing_Transformer.ipynb # How the data passed into the transformer was pre-processed, generated output stored in /transformer_data/
└── transformer_data/ # Model training notebooks
│ ├── labels.pt # Labels tensor
│ ├── other_features_file.pt # Patient demographic tensor
│ └── output_file.pt # Pre-processed ECG tensor NOTE: PLEASE DOWNLOAD FROM GOOGLE DRIVE
│
├── data/ # Data files
│ └── coorteeqsrafva.csv # Patient demographic CSV
│ └── ecgeq-500hzsrfava.npy # ECG waveform data for all patients NOTE: PLEASE DOWNLOAD FROM GOOGLE DRIVE
│ └── af_dataset.csv # Combined ECG waveform and patient demographic data NOTE: PLEASE DOWNLOAD FROM GOOGLE DRIVE
│
├── saved_models/ # Saved Models
│ └── simple_mlp_model.pth # MLP saved model
│ └── lstm_saved.h5 # LSTM saved model
│ └── logistic_regression_model.pkl # Logistic Regression saved model
│ └── transformer_model_unbalanced.pth # Unweighted transformer saved model
│ └── transformer_model.pth # Weighted transformer saved model
│
└── README.md # Project documentation
The main file to run is Inference Notebok.ipynb as it utilizes all the saved models and then runs inference with these models. Please ensure the file structure follows the convention as outlined in File Structure.
The rest of the notebooks under src contain the code for the preprocessing and the training of the models that can be looked at as a reference.