dumplingman0403/ECG-GAN

Synthesize plausible ECG signals via Generative adversarial networks

Synthesis ECG signals via Generative adversarial networks

Setup

pip install -r requirement.txt

Usage

See data description

download dataset

• option 1

sh download.sh

if you haven't install unzip, install unzip package first before run download.sh
install guildline:
General - for Linux
Homebrew - for MacOS

• option 2: download from website
  • AF Classification from a Short Single Lead ECG Recording - The PhysioNet Computing in Cardiology Challenge 2017
  • MIT-BIH Arrhythmia Database

Process ECG signals

run following command

python3 process_ecg.py

if you change the dataset path, modify process_ecg.py

# modify dataset path if necessary
AA_DATASET_DIR = 'AA_dataset/'   # MIT-BIH Arrhythmia Database
    
AF_DATASET_DIR = 'AF_dataset/'   # AF Classification from a Short Single Lead ECG Recording - The PhysioNet Computing in Cardiology Challenge 2017
LABEL_PATH = 'AF_dataset/REFERENCE-original.csv'

GAN Training

training ECG signal with GAN model

python3 train.py

modify the input dataset path in train.py if you change the path of X_train_af.pkl and y_af.pkl

X_train = pickle.load(open("path_of_X_train_af.pkl", "rb"))  # --> load AF dataset
y = pickle.load(open("path_of_y_af.pkl", 'rb'))

Output

MIT-BIH Arrhythmia Database

Short Single Lead ECG Recording

Atrial Fibrillation

real

epoch 1000

epoch 2000