https://github.com/hase3b/bonbid-hie-segmentation
This repository implements a 3D U-Net model for segmenting HIE lesions in neonatal MRI scans, exploring various loss functions including Dice, Dice-Focal, Tversky, Hausdorff Distance, and hybrid loss functions. BONBID-HIE dataset is used for this study.
https://github.com/hase3b/bonbid-hie-segmentation
lesions loss-functions miccai2023 monai multimodality pytorch segmentation simpleitk torchio unet
Last synced: 8 months ago
JSON representation
This repository implements a 3D U-Net model for segmenting HIE lesions in neonatal MRI scans, exploring various loss functions including Dice, Dice-Focal, Tversky, Hausdorff Distance, and hybrid loss functions. BONBID-HIE dataset is used for this study.
- Host: GitHub
- URL: https://github.com/hase3b/bonbid-hie-segmentation
- Owner: hase3b
- Created: 2025-01-08T10:05:47.000Z (10 months ago)
- Default Branch: main
- Last Pushed: 2025-02-06T16:43:33.000Z (9 months ago)
- Last Synced: 2025-02-06T17:39:38.816Z (9 months ago)
- Topics: lesions, loss-functions, miccai2023, monai, multimodality, pytorch, segmentation, simpleitk, torchio, unet
- Language: Jupyter Notebook
- Homepage:
- Size: 3.39 MB
- Stars: 0
- Watchers: 1
- Forks: 0
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
Awesome Lists containing this project
README
# Multimodal HIE Lesion Segmentation in Neonates: A Study of Loss Functions
This repository implements a 3D U-Net for the segmentation of Hypoxic-Ischemic Encephalopathy (HIE) lesions in neonatal MRI scans. The project utilizes the BONBID-HIE dataset and explores the optimization of loss functions for improved segmentation accuracy. Specifically, we investigate several loss functions including Dice Loss, Dice-Focal Loss, Tversky Loss, and Hausdorff Distance Loss, along with two novel hybrid loss functions that combine Dice-Focal and Tversky Loss with Hausdorff Distance Loss. Our findings highlight the superior performance of the Tversky-HausdorffDT Loss.
This is a term project for Deep Learning (96943) Fall 24' course taught by Dr. Tahir Syed at the Institute of Business Administration (IBA) Karachi. It's a collaborative effort of Abdul Haseeb and Annayah Usman.
## **Workflow**
*Figure 1: Brief Overview of the Project*
## **Visualizing The Problem**
*Figure 2: Sample 3D ADC Map*
*Figure 3: Sample 3D ZADC Map*
*Figure 4: Sample 3D Binary Label Mask*
*Figure 5: Ground Truth Imposed on ADC and ZADC Maps Across Axial Slices*
## **Data and Preprocessing**
BONBID-HIE MICCAI 2023 Challenge dataset is used, consisting of 3D Apparent Diffusion Coefficient (ADC) maps, Z-score normalized ADC maps (ZADC), and binary label masks for 133 HIE patients.
* Preprocessing Steps:
* Resampling the ADC and ZADC maps to a fixed size of (192, 192, 32).
* Intensity normalization of the resampled maps using mean and standard deviation.
* Concatenation of the ADC and ZADC maps to create a 2-channel input for the U-Net model.
* Data Augmentation:
* Random Noise
* Random Anisotropy
* Random Blur
* Random Gamma
* Random Elastic Transformation
These augmentations simulate real-world imaging variations and improve the model's generalizability. More details on this in the "Term Paper.PDF".
## **Model Architecture**
A 3D U-Net architecture is employed due to its strong inductive bias, especially when dealing with small datasets. The model consists of:
* Three encoder and decoder blocks.
* Each encoder block includes double 3D convolutional layers, batch normalization, LeakyReLU activation, dropout, and max-pooling.
* The decoder block progressively upsamples using transpose convolutions, integrating encoder features via skip connections.
The model uses batch normalization for faster convergence and dropout to reduce overfitting. LeakyReLU ensures gradient flow through non-lesion areas, which is critical for sparse lesion segmentation.
## **Loss Functions**
We explore and compare the following loss functions:
* Dice Loss: Measures the overlap between predicted and ground truth masks.
* Dice-Focal Loss: Combines Dice Loss with Focal Loss to prioritize hard-to-classify examples.
* Tversky Loss: Generalizes Dice Loss to handle class imbalances, particularly useful for medical segmentation.
* Hausdorff Distance Loss (HDTL): Measures the distance between predicted and true boundaries for more precise segmentation.
* Hybrid Loss Functions:
* DiceFocal-HausdorffDT Loss: Combines Dice-Focal Loss with Hausdorff Distance Loss.
* Tversky-HausdorffDT Loss: Combines Tversky Loss with Hausdorff Distance Loss.
## **Evaluation Metrics**
Segmentation performance is evaluated using the following metrics:
* Dice Coefficient: Measures the volumetric overlap between the predicted and ground truth masks.
* Mean Surface Distance (MSD): Measures the average distance between the surfaces of predicted and true masks.
* Normalized Surface Dice (NSD): Evaluates the similarity between the boundary surfaces of predicted and ground truth masks.
## **Results**
| Loss Function | Dice ↑ | MSD ↓ | NSD ↑ | Epochs |
|-----------------------------|--------|---------|--------|--------|
| Dice Loss (Baseline) | 0.3800 | 15.0650 | 0.3850 | 32 |
| Dice Focal Loss | 0.4900 | 1.7925 | 0.5275 | 49 |
| Tversky Loss | 0.3525 | 15.3650 | 0.3375 | 38 |
| HausdorffDT Loss | 0.3300 | Inf | 0.2800 | 29 |
| DiceFocal-HausdorffDT Loss | 0.4925 | 1.4225 | 0.5300 | 72 |
| Tversky-HausdorffDT Loss | 0.5000 | 1.6250 | 0.5325 | 59 |
*Table 1: Metric-Based Comparison of Loss Functions*
*Figure 6: Visualizing Segmentation Masks Across Loss Functions*
* Hybrid loss functions perform better than standalone loss function but take longer to converge.
* Tversky-HausdorffDT Loss achieved the highest performance in terms of Dice and NSD while maintaining competitive MSD.
## **Limitations**
* Small lesions (< 1% volume) are difficult to segment accurately due to dataset heterogeneity.
* Evaluation was limited to the validation set, and the performance on the test set remains unknown.
* Resampling of label masks could introduce distortion in segmentation regions.
## **Future Work**
* Defining a custom loss function specifically for HIE lesions.
* Exploring the use of zero-shot models like MedSAM-2 for improved segmentation without the need for extensive training data.
To get a comprehensive understanding of the project, please refer to the "Term Paper.PDF" in the repo. Refer to "Preprocessing.ipynb" for code related to preprocessing and "Main.ipynb" for the rest of the code.
## **Repository Contents**
| Folder/File | Description |
| ---------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------- |
| /Data/ | Raw and preprocessed data files aren't uploaded due to GitHub's upload constraints. However, it contains notes regarding the preprocessed dataset reproducibility. |
| /Preprocessing.ipynb | Notebook to convert raw files to preprocessed maps and labels through resampling, normalization, and concatenation. |
| /Main.ipynb | Main project notebook containing data loading, augmentation, training, evaluation, and experimentation process. |
| /Term Paper.PDF | Detailed paper explaining the project, including defense of different strategies used and discussion on findings. |
## **Acknowledgments**
* Tools: NumPy, Matplotlib, ImageIO, Torch, TorchIO, SimpleITK, MONAI (Metrics, Transforms, Losses), Pandas, TorchSummary
* Model(s): 3D U-Net
* Data Source: [](https://doi.org/10.5281/zenodo.10602767)
* Instructor: Dr Tahir Syed (Professor IBA Karachi)