Hypernet-Ensemble Learning of Segmentation Probability for Medical Image Segmentation with Ambiguous Labels
The implementation of Hypernet-Ensemble Learning of Segmentation Probability for Medical Image Segmentation with Ambiguous Labels.
We aim to improve the segmentation probability estimation of DL networks for medical image segmentation that often suffers from the ambiguity in human annotations and anatomical/pathological structures using the hypernetwork ensemble strategy with the varying Tversky loss.
Hypernetwork Architecture
The overall hypernetwork architecture shown above is implemented in networks/hyper_resunet.py.
The hyperconvolution blocks are implemented in blocks/hyper_convolution.py.
Usage
The codes are implemented with the MONAI framework (https://monai.io/), PyTorch (https://pytorch.org/), and PyTorch Lightning (https://www.pytorchlightning.ai/).
Requirements
pip install -r requirements.txt
Example Dependencies
- monai=0.7.0
- torch=1.9.1
- nibabel=3.2.1
Training
The hypernetwork and optimizer is wrapped with PyTorch Lightning (lightning_modules/module_hyper_resunet.py).
Please see train.py for setting up the network parameters and training configurations.
Inference
Please see predict.py for inference.
The network parameters need to be the same with a trained network.
Contact
Please contact Sungmin Hong (MGH, HMS [email protected]) if you have questions on the codes or the paper.