Electrocardiogram Generation using Diffusion

A latent diffusion model for generating ECG signals

Getting started

Clone the repository
Install the dependencies with pip install -r requirements.txt
Copy the .env.example file in the project root and rename it to .env
Add your environment variables to the file to be able to use some of the features
Using the code
1. Run the train.py file for training models
2. Run the test.py file to generate synthetic data

Dataset

Download the PTB.zip file on Wetransfer
Place the unzipped file inside the root directory of the project
Run the data_prepare.py script
1. Based on the configuration you may need to modify the directory paths found in the cofigurations.py script
2. After the script is done, the data should be split in train, test, and validation directories and each .csv file should have a corresponding .pt file right below it
The extracted and normalized dataset should be in the project root directory

Trained Models

Install DVC on your machine if it is not already installed
Run the dvc pull command on the project root to get all the latest model files
You may need to authenticate with your Google account in order to have access
Now you can run the test.py script to generate data from the pre-trained models

Configurations

All the configuration variables live in the configurations.py file. There you can switch between models to train and sample. If you are going to do your own training the hyperparameters of the networks can be found there as well.

Training
1. Choose a model to train in the configurations.py script (the options are commented out)
2. Set up WandB credentials if you want artifacts and media to be reported there
3. Run the train.py script, models will be saved each 10th epoch
Sampling
1. Choose a model to sample from in the configurations.py script (the options are commented out)
2. Run the test.py script, generated files and plots will be saved in the /data directory

Collaborators