This is the official code for the paper "An Elastic Interaction-Based Loss Function for Medical Image Segmentation" presented at MICCAI 2020 (https://arxiv.org/abs/2007.02663).
- Updated the code to Pytorch 1.11.0
- Fixed some bugs and improved performance
- Python 3.7 or higher
- Pytorch 1.11.0 or higher
To train your model, run the following command:
python train.py --train_dataset '/your_training_data_path' --test_dataset '/your_test_data' --save_path '/save_model_path'
You can also specify other arguments such as batch size, learning rate, number of epochs, etc. See train.py
for more details.
If you want to customize your Dataset: modify the ImageToImage2D
in ./unet/dataset.py
.
The elastic interaction loss file is located in ./unet
folder. You can import it and use it as a custom loss function for your segmentation model.
If you find this code useful, please cite our paper:
@inproceedings{LanXZ20,
author = {Yuan Lan and
Yang Xiang and
Luchan Zhang},
title = {An Elastic Interaction-Based Loss Function for Medical Image Segmentation},
booktitle = {Medical Image Computing and Computer Assisted Intervention - {MICCAI}
2020 - 23rd International Conference, Lima, Peru, October 4-8, 2020,
Proceedings, Part {V}},
series = {Lecture Notes in Computer Science},
volume = {12265},
pages = {755--764},
publisher = {Springer},
year = {2020}
}
This code is based on the Pytorch UNet template from https://github.com/cosmic-cortex/pytorch-UNet. We thank the authors for their work.