3D image segmentation with UNet and spatial data augmentation
- Free software: GNU General Public License v3
This is an implementation of 3D-UNet with spatial data augmentation. The 3D UNet takes the 3D image and segments the object.
It can perform binary segmenation, or even multi label segmentation. The full 3D deformation makes this model robust in low supervised setting.
To install this package,
-
cd ./3D_UNet_segmentaion
-
pip install -e .
pip install -e . will install required dependencies. I recommend to create new conda environment before install this package.
To train the model, pass the training data directory to ./segmentation_3D/train.py
The input image should be in [z,y,x,channel] format.
run train.py by passing required arguments like following Example
Example: python train.py --source /path/to/source/directory --target /path/to/target/directory --data_type tiff --ch 2 --save_name demo --iters 100 --save_iters 10 --batch 20
To inference, pass the test image directory and other arugmentsto ./segmentation_3D/test.py
The input image should be in [z,y,x,channel] format.
The segmented image will be saved to ./inference/save_dir
Example: python test.py --dir /path/to/test/image/directory --data_type tiff --ch 2 --model demo --iters 100 --save_dir directory name to save result
This 3D UNet with spatial data augmentation used in segmentation model comparison in low supervised setting:
"Few Labeled Atlases are Necessary for Deep-Learning-Based Segmentation."
Hyeon Woo Lee, Mert R. Sabuncu, and Adrian V. Dalca.
arXiv preprint arXiv:1908.04466 (2019).
- Our proposed method in this paper, multi-atlas segmentation with semi-supervision can be found from: voxelmorph.mit.edu
Spatial Transformer Network implementation code is from:
Unsupervised Learning for Fast Probabilistic Diffeomorphic Registration
Adrian V. Dalca, Guha Balakrishnan, John Guttag, Mert R. Sabuncu
MICCAI 2018. eprint arXiv:1805.04605
For and problems or questions, please send me an email at [email protected]
This package was created with Cookiecutter and the audreyr/cookiecutter-pypackage project template.