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 (https://arxiv.org/abs/2112.06693).
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.
The overall hypernetwork architecture shown above is implemented in networks/hyper_resunet.py.
The hyperconvolution blocks are implemented in blocks/hyper_convolution.py.
The codes are implemented with the MONAI framework (https://monai.io/), PyTorch (https://pytorch.org/), and PyTorch Lightning (https://www.pytorchlightning.ai/).
pip install -r requirements.txt
Example Dependencies
- monai=0.7.0
- torch=1.9.1
- nibabel=3.2.1
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.
Please see predict.py for inference.
The network parameters need to be the same with a trained network.
@misc{hong2021hypernetensemble,
title={Hypernet-Ensemble Learning of Segmentation Probability for Medical Image Segmentation with Ambiguous Labels},
author={Sungmin Hong and Anna K. Bonkhoff and Andrew Hoopes and Martin Bretzner and Markus D. Schirmer and Anne-Katrin Giese and Adrian V. Dalca and Polina Golland and Natalia S. Rost},
year={2021},
eprint={2112.06693},
archivePrefix={arXiv},
primaryClass={eess.IV}
}
Please contact Sungmin Hong (MGH, HMS shong20@mgh.harvard.edu) if you have questions on the codes or the paper.
Update March 2022: I joined AWS ML Solution Lab. My new email address is hsungmin@amazon.com.