CardioNet: Deep Learning Framework for ECG Arrhythmia Detection

Complete PyTorch project for classifying ECG heartbeats from the MIT-BIH heartbeat CSV dataset.

Project Structure

.
├── data/
│   └── MIT-BIH.csv
├── dashboard/
│   └── app.py
├── results/
├── src/
│   ├── dataset.py
│   └── models.py
├── train.py
├── evaluate.py
├── requirements.txt
└── README.md

Dataset Format

Place the MIT-BIH heartbeat CSV file in the data/ folder:

• data/MIT-BIH.csv

Each row must contain:

• Optional first column named Unnamed: 0, which is dropped automatically
• 180 ECG signal values in columns 0 to 179
• 1 label column named target
• 5 classes labeled 0 to 4

Class mapping used in the project:

Label	Class
0	Normal
1	Supraventricular premature
2	Premature ventricular contraction
3	Fusion of ventricular and normal
4	Unclassifiable

Split Strategy

The single dataset is split once with stratified sampling:

• 70% training set: used to fit model weights and compute normalization mean/std.
• 15% validation set: used during training to choose the best checkpoint by validation accuracy.
• 15% test set: kept untouched until evaluate.py computes final metrics.

All splits preserve class proportions as closely as possible and use random_state=42.

Installation

Create and activate a virtual environment, then install dependencies:

pip install -r requirements.txt

numpy is pinned below version 2 to avoid binary compatibility issues with older local Matplotlib builds.

Train the CNN

python train.py --model cnn --epochs 20 --batch-size 128

Running python train.py without arguments performs a quick one-epoch CNN run; increase --epochs for stronger final results.

The training script:

• Loads data/MIT-BIH.csv
• Creates a stratified 70/15/15 train-validation-test split with random_state=42
• Trains only on the 70% training split
• Validates only on the 15% validation split
• Leaves the 15% test split untouched for final evaluation
• Normalizes ECG values using training statistics
• Applies optional training augmentation
• Trains with CrossEntropyLoss and Adam
• Saves the best CNN checkpoint to results/best_cnn_model.pth
• Saves learning curves to results/cnn_training_curves.png

Disable data augmentation if needed:

python train.py --model cnn --no-augment

Train the MLP Baseline

python train.py --model mlp --epochs 20 --batch-size 128

This saves the best MLP checkpoint to results/best_mlp_model.pth.

Evaluate the CNN

python evaluate.py

The evaluation script saves:

• results/metrics.json
• results/confusion_matrix.png

Evaluation uses only the untouched 15% test split. Metrics include accuracy, weighted precision, weighted recall, weighted F1-score, a full classification report, the confusion matrix, and split information.

Run the Streamlit Dashboard

streamlit run dashboard/app.py

The dashboard can:

• Load the trained CNN checkpoint
• Select a local sample from data/MIT-BIH.csv
• Plot the ECG signal
• Predict the heartbeat class
• Show confidence scores
• Display saved metrics and the confusion matrix from the untouched 15% test split

Notes

• CNN inputs are reshaped to [batch_size, 1, 180].
• MLP inputs remain flattened as [batch_size, 180].
• The model checkpoint stores normalization mean and standard deviation, so training, evaluation, and dashboard prediction use consistent preprocessing.
• If CUDA is available, training and evaluation automatically use the GPU.
• The split is a random beat-level split. If multiple beats from the same patient are present, this may still cause patient-level leakage; a patient-level split would be stricter for clinical evaluation.