NicolasNoya/Heart_Segmentation

Cardiac Pathology Prediction - DenseNet

Model Description

This is a custom-implemented DenseNet for cardiac MRI segmentation and pathology prediction. The model performs semantic segmentation of cardiac structures (left ventricle, right ventricle, and myocardium) from short-axis cardiac cine MR images.

Key Features

• 🏗️ Fully Convolutional DenseNet architecture implemented from scratch
• 🫀 4-class semantic segmentation (background, LV, RV, myocardium)
• 🔬 Trained on NIfTI format cardiac MRI data
• 📊 Combined Cross-Entropy and Dice Loss
• 🎯 Designed for cardiac pathology classification

Architecture

The model follows a U-Net style encoder-decoder architecture with dense blocks:

• Input: Single-channel 2D cardiac MRI slices (128×128)
• Encoder: 3 dense blocks (3, 4, 5 layers) with transition down
• Bottleneck: Dense block with 8 layers
• Decoder: 3 dense blocks (5, 4, 3 layers) with skip connections
• Output: 4-channel probability map (softmax activated)
• Growth rate: 8 (bottleneck: 7)

Intended Use

This model is designed for:

• Cardiac MRI segmentation research
• Educational purposes in medical image analysis
• Baseline comparison for cardiac segmentation tasks
• Feature extraction for cardiac pathology classification

Note: This model is for research purposes only and not intended for clinical use.

Training Details

• Loss Function: Combined Cross-Entropy and Dice Loss (α=0.25)
• Optimizer: Adam
• Framework: PyTorch 2.6.0
• Data Format: NIfTI (.nii) files
• Image Size: 128×128 pixels
• Preprocessing: ROI extraction, normalization, data augmentation

How to Use

import torch
from densenet.densenet import DenseNet

# Load the model
model = DenseNet()
model.load_state_dict(torch.load('model_weights.pth', map_location='cpu'))
model.eval()

# Inference
with torch.no_grad():
    # input_image: (1, 1, 128, 128) tensor
    output = model(input_image)  # (1, 4, 128, 128)
    prediction = torch.argmax(output, dim=1)  # Get class predictions

Model Files

• model_weights.pth: Complete model checkpoint including weights and optimizer state
• densenet/: Source code for the DenseNet architecture
• Full implementation available at: GitHub Repository

Citation

If you use this model, please cite the original paper that inspired this implementation:

Densely Connected Fully Convolutional Network for Short-Axis Cardiac Cine MR Image 
Segmentation and Heart Diagnosis Using Random Forest

Author

Francisco Nicolás Noya

• GitHub: @NicolasNoya
• Project: IMA205 Challenge

License

MIT License - See repository for details.

Disclaimer

This model is provided for research and educational purposes only. It has not been clinically validated and should not be used for medical diagnosis or treatment decisions.