Deep Learning-based cardiac MRI segmentation and pathology classification using a custom-implemented DenseNet architecture for automated cardiovascular disease diagnosis.
- 🤗 Pre-trained Model: Download from Hugging Face
- 📊 Model Weights: Ready-to-use trained weights available
This project implements an end-to-end deep learning pipeline for cardiac pathology prediction from short-axis cardiac cine MR images. The system performs semantic segmentation of cardiac structures (left ventricle, right ventricle, and myocardium) using a fully convolutional DenseNet, followed by feature extraction and pathology classification.
Key Features:
- 🏗️ Custom DenseNet implementation built from scratch in PyTorch
- 🔬 Medical image processing pipeline for NIfTI (.nii) format MRI data
- 🎯 Multi-class segmentation of cardiac structures (4 classes)
- 📊 Automated feature extraction from segmented cardiac regions
- 🧠 Random Forest classifier for pathology prediction
- 📈 Comprehensive training pipeline with TensorBoard profiling
- 🔄 Data augmentation and Region of Interest (ROI) extraction
The project is based on the paper: "Densely Connected Fully Convolutional Network for Short-Axis Cardiac Cine MR Image Segmentation and Heart Diagnosis Using Random Forest"
CardiacPathologyPrediction/
├── densenet/ # Custom DenseNet implementation
│ ├── densenet.py # Main network architecture
│ ├── dense_block.py # Dense block components
│ ├── layer.py # Individual layer definitions
│ ├── transitions.py # Transition layers (up/down sampling)
│ └── custom_loss.py # Combined loss function (CE + Dice)
├── research/ # Experimental scripts and prototypes
├── data/ # Dataset directory (not included in repo)
│ ├── Train/ # Training MRI images
│ ├── Test/ # Test MRI images
│ ├── metaDataTrain.csv # Patient metadata
│ └── metaDataTest.csv # Test metadata
├── runs/ # TensorBoard logs
├── densnet_trainer.py # Training pipeline orchestration
├── feature_extractor.py # Cardiac feature computation
├── niidataloader.py # Custom PyTorch dataset for NIfTI files
├── roi.py # Region of Interest extraction
├── profiler.py # TensorBoard profiling utilities
├── model_weights.pth # Trained model checkpoint
├── presentation.ipynb # Results and analysis notebook
└── requirements.txt # Python dependencies
- Python 3.10+
- CUDA-compatible GPU (recommended for training)
- 8GB+ RAM
- Clone the repository
git clone https://github.com/NicolasNoya/CardiacPathologyPrediction.git
cd CardiacPathologyPrediction- Create a virtual environment
python -m venv .venv
source .venv/bin/activate # On Windows: .venv\Scripts\activate-
Install PyTorch (select the appropriate version for your system)
Visit PyTorch Get Started and choose your configuration.
Example for CUDA 11.8:
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118
-
Install dependencies
pip install -r requirements.txtfrom densnet_trainer import DenseNetTrainer
# Initialize trainer
trainer = DenseNetTrainer(
path_to_images="data/Train",
epochs=100,
alpha=0.25, # Combined loss weight
train_fraction=0.8,
check_val_every=10
)
# Train the model
trainer.train()Option 1: From Hugging Face (Recommended)
from huggingface_hub import hf_hub_download
from densenet.densenet import DenseNet
import torch
# Download weights from Hugging Face
model_path = hf_hub_download(
repo_id="NicolasNoya/Heart_Segmentation",
filename="model_weights.pth"
)
# Load model
model = DenseNet()
model.load_state_dict(torch.load(model_path, map_location='cpu'))
model.eval()Option 2: From Local File
from densenet.densenet import DenseNet
import torch
model = DenseNet()
model.load_model("model_weights.pth")
model.eval()from feature_extractor import FeatureExtractor
extractor = FeatureExtractor()
features = extractor.extract_features(
diastole_image,
systole_image,
height,
weight,
voxel_spacing
)tensorboard --logdir=runsThen navigate to http://localhost:6006 in your browser.
The DenseNet architecture consists of:
- Inception-X initial layer (1 → 24 channels)
- Encoder path: 3 dense blocks (3, 4, 5 layers) with transition down
- Bottleneck: Dense block with 8 layers
- Decoder path: 3 dense blocks (5, 4, 3 layers) with transition up and skip connections
- Output: 1×1 convolution + softmax (4 classes)
- Growth rate: 8 (bottleneck: 7)
Loss Function: Combined Cross-Entropy and Dice Loss
L = α · L_CE + (1 - α) · L_DiceThe model achieves competitive performance on cardiac MRI segmentation:
- Effective segmentation of left ventricle, right ventricle, and myocardium
- Robust to variations in cardiac anatomy
- Efficient feature extraction for downstream pathology classification
Detailed results and visualizations are available in presentation.ipynb.
To make your model publicly accessible and showcase it to potential employers, you can upload it to Hugging Face Hub:
- Create an account at huggingface.co
- Go to Settings → Access Tokens
- Create a new token with write permissions
- Copy the token
# Set your token as environment variable (recommended)
export HF_TOKEN=your_token_here
# Run the upload script
python upload_to_huggingface.py --repo-name your-username/cardiac-densenetOr pass the token directly:
python upload_to_huggingface.py --token your_token_here --repo-name your-username/cardiac-densenetOnce uploaded, anyone can download and use your model:
from huggingface_hub import hf_hub_download
import torch
# Download weights
model_path = hf_hub_download(
repo_id="your-username/cardiac-densenet",
filename="model_weights.pth"
)
# Load and use
model = DenseNet()
model.load_state_dict(torch.load(model_path))
model.eval()Benefits:
- 🌐 Publicly accessible model weights
- 📊 Automatic model card generation
- 🔍 Discoverable in Hugging Face search
- 💼 Great for portfolio and job applications
- 📈 Track downloads and usage
| Module | Description |
|---|---|
densnet_trainer.py |
Complete training pipeline with validation, metrics, and checkpointing |
feature_extractor.py |
Extracts clinical features (volumes, ejection fraction, mass, etc.) |
niidataloader.py |
PyTorch Dataset for loading NIfTI medical images with voxel spacing |
roi.py |
Region of Interest extraction for focused cardiac analysis |
profiler.py |
TensorBoard integration for training visualization and analysis |
All DenseNet components were implemented from scratch:
- Dense blocks with skip connections
- Transition layers (downsampling/upsampling)
- Custom loss functions
- Inception-style initial layer
- Medical Imaging: Custom dataloader for NIfTI format with proper voxel spacing handling
- Deep Learning: Implemented DenseNet architecture with skip connections and dense blocks
- Computer Vision: Semantic segmentation with connected component analysis
- Data Augmentation: Rotation, flipping, and intensity transformations
- Evaluation Metrics: Dice coefficient, IoU, confusion matrix
- Profiling: TensorBoard integration for comprehensive training analysis
- GPU Recommended: Training and inference are computationally intensive. A CUDA-enabled GPU is strongly recommended.
- Dataset: The dataset is not included in this repository due to size and privacy constraints. The data structure expects NIfTI (.nii) format files organized by patient ID.
- Research Folder: Contains experimental code and prototype implementations used during development.
This project was developed as part of the IMA205 course challenge. Contributions, issues, and feature requests are welcome!
Francisco Nicolás Noya
- GitHub: @NicolasNoya
- Project: IMA205 Challenge - Cardiac Pathology Prediction
This project is licensed under the MIT License - see the LICENSE file for details.
- Based on research from: "Densely Connected Fully Convolutional Network for Short-Axis Cardiac Cine MR Image Segmentation and Heart Diagnosis Using Random Forest"
- Built with PyTorch, TensorFlow, and scikit-learn
- Medical imaging processing with NiBabel
⭐ If you find this project useful, please consider giving it a star!
For questions or collaborations, feel free to reach out through GitHub issues or discussions.