CardioAI: Heart Disease Detection Across 4 Datasets

💡 Inspiration

Cardiovascular disease (CVD) is the #1 cause of death globally, killing over 18 million people every year. Most deaths are preventable with early detection. We wanted to build an AI system that could predict heart disease risk across multiple real-world biomedical datasets — not just one — and make every prediction explainable to doctors.

🔨 What We Built

A complete end-to-end machine learning pipeline that:

• Trains and evaluates 5 ML models on 3 structured cardiac datasets
• Trains a 1D Convolutional Neural Network on raw ECG timeseries signals
• Uses Optuna (60-trial Bayesian optimization) to tune hyperparameters
• Generates SHAP explanations for every prediction (feature importance, beeswarm, waterfall, dependence plots)

📁 Datasets Used

⚙️ How We Built It

1. Data Preprocessing — RobustScaler, median imputation, SMOTE for class imbalance, LabelEncoder for categorical features
2. Baseline Models — Logistic Regression, Random Forest, XGBoost, LightGBM, CatBoost (all benchmarked with Accuracy, F1, ROC-AUC)
3. Hyperparameter Tuning — Optuna TPE Sampler, 60 trials, 5-fold Stratified Cross-Validation
4. Deep Learning — 5-layer 1D-CNN with BatchNorm, AdaptiveAvgPool, Dropout, CosineAnnealingLR scheduler, class-weighted CrossEntropyLoss
5. Explainability — SHAP TreeExplainer for all tree models with global + per-patient explanations

📊 Results

🏆 Best ROC-AUC: 0.9412 on Heart Attack dataset
🏆 Best Accuracy: 0.9438 on ECG Timeseries (1D-CNN)

🧠 What We Learned

• How to handle severely imbalanced medical datasets using SMOTE
• How Optuna Bayesian optimization outperforms random/grid search
• How SHAP values make black-box ML models clinically trustworthy
• How 1D-CNNs can extract patterns from raw ECG waveforms without manual feature engineering

🚧 Challenges We Faced

• The Cardio Base CSV used a semicolon separator instead of a comma — caused all columns to merge into one string
• The ECG heatmap (180×180 cells) froze Colab for 20+ minutes — solved by computing only top-20 correlated features
• Class imbalance across datasets required different SMOTE thresholds per dataset
• Keeping the pipeline fully modular so any new cardiac dataset can be plugged in with minimal changes

🔮 Future Work

• Deploy as a clinical web app using Streamlit or Flask
• Add LSTM / Transformer for ECG timeseries
• Integrate federated learning to train on hospital data privately
• Add real-time ECG device input via wearable APIs

Built With

• catboost
• google
• google-colab
• imbalanced-learn-(smote)
• lightgbm
• matplotlib
• numpy
• optuna
• pandas
• python
• pytorch
• scikit-learn
• seaborn
• shap
• xgboost