3D-UCaps: 3D Capsules Unet for Volumetric Image Segmentation (MICCAI 2021)
VinAIResearch, updated
🕥
2022-01-21 05:54:13
3D-UCaps: 3D Capsules Unet for Volumetric Image Segmentation
3D-UCaps is a voxel-based Capsule network for medical image segmentation. Our architecture is based on the symmetry U-net with two parts: the encoder forms by Capsule layers, whereas the decoder contains traditional convolutional layers. 3D-UCaps, therefore inherits the merits from both Capsule networks to preserve the part-to-whole relationship and CNNs to learn translational invariant representation. We conducted experiments on various datasets (including iSeg-2017, LUNA16, Hippocampus, and Cardiac) to demonstrate the superior performance of 3D-UCaps, where our method outperforms the baseline method SegCaps while being more robust against rotational transformation when compared to 3D-Unet.
Details of the UCaps model architecture and experimental results can be found in our following paper:
@inproceedings{nguyen20213d,
title={3D-UCaps: 3D Capsules Unet for Volumetric Image Segmentation},
author={Nguyen, Tan and Hua, Binh-Son and Le, Ngan},
booktitle={International Conference on Medical Image Computing and Computer-Assisted Intervention},
pages={548--558},
year={2021},
organization={Springer}
}
Please CITE our paper when UCaps is used to help produce published results or incorporated into other software
Usage
Installation
We provide instructions on how to install dependencies via conda. First, clone the repository locally:
git clone https://github.com/VinAIResearch/3D-UCaps.git
Then, install dependencies depends on your cuda version. We provide two versions for CUDA 10 and CUDA 11
conda env create -f environment_cuda11.yml
or
conda env create -f environment_cuda10.yml
Data preparation
Download and extract these datasets: * iSeg-2017 challenge (infant brain MRI segmentation): https://iseg2017.web.unc.edu/download/ * Lung Nodule Analysis 2016 (LUNA16): https://luna16.grand-challenge.org/Download/ * Cardiac and Hippocampus dataset from Medical Segmentation Decathlon: http://medicaldecathlon.com/
We expect the directory structure to be the following: ``` path/to/iseg/ domainA/ domainA_val/
path/to/cardiac/ imagesTr labelsTr
path/to/hippocampus/ imagesTr labelsTr
path/to/luna/ imgs segs ```
Note: there are some files in LUNA16 dataset can lead to an error when training so we have removed it:
1.3.6.1.4.1.14519.5.2.1.6279.6001.771741891125176943862272696845.mhd
1.3.6.1.4.1.14519.5.2.1.6279.6001.927394449308471452920270961822.mhd
Training
Arguments for training can be divided into 3 groups:
-
Trainer args to initialize Trainer class from Pytorch Lightning.
-
Important arguments:
gpus,accelerator,check_val_every_n_epoch,max_epochs. - Fix arguments in
train.py:benchmark,logger,callbacks,num_sanity_val_steps,terminate_on_nan -
Model args depend on which model you use (UCaps, SegCaps or U-net) and defined in
add_model_specific_argsmethod of that module. -
Important arguments:
in_channels,out_channels,val_frequency,val_patch_size,sw_batch_size,overlap. The last three args are use in sliding window inference method from MONAI library. -
Args specific for training:
root_dir,log_dir,dataset,fold,cache_rate,cache_dir,model_name,train_patch_size,num_workers,batch_size,num_samples. -
cache_rateandcache_dirdefine whether you want to use CacheDataset or PersistentDataset when loading data. num_samplesis a arg in RandCropByPosNegLabel method, the effective batch size isbatch_sizexnum_samples.
The full list of arguments can be shown through the command:
python train.py -h
We provide bash script with our config to train UCaps model on all datasets and can be run as follow:
bash scripts/train_ucaps_iseg.sh
Validation
Arguments for validation can be divided into 3 groups:
- Trainer args to initialize Trainer class. The only argument we need to use here is
gpus. - Args for sliding window inference method
- Args specific for validation
root_dir,output_dir,save_image,model_name,dataset,fold,checkpoint_path
The full list of arguments can be shown through the command:
python evaluate.py -h
We provide bash script with our config to validate trained UCaps models on all datasets, you just need to download our models in Model Zoo and put them in logs folder. After that, you can run the evaluation script for targeted dataset as follow:
bash scripts/evaluate_ucaps_iseg.sh
Rotation experiment
Same with validation but add two more arguments rotate_angle (in degree) and axis (z/y/x or all) to create test rotated subject.
The full list of arguments can be shown through the command:
python evaluate_iseg.py -h
We provide bash script with our config to compare between trained UCaps (download) and U-net (download) on subject 9th of iSeg-2017 dataset, the first arugment is rotate_angle and the second argument is axis:
bash scripts/evaluate_rotation.sh 0 z
Rotation experiment on SkipDenseSeg model
- Cloning SkipDenseSeg project
- Replacing their
val.pywith ourval.py - Running
val.pywith args, for example:
python val.py --gpu 1 --sw_batch_size 32 --overlap 0.75 --output_dir=/home/ubuntu/
Model Zoo
About the code This repository has been refactored to use Pytorch Lightning framework and MONAI library for data preprocessing, data loading, inferencing to ensure the reproducibility and extendability of our work as well as improve efficiency when training. Hence, the results here have been improved a little bit when compared to their counterparts in the paper.
Dice Coefficient on subject 9th of iSeg-2017 dataset:
| Model | CSF | GM | WM | Average | Pretrained model | |-------|:---:|:---:|:---:|:-----:|------------------| | 3D-UCaps | 95.01 | 91.51 | 90.59 | 92.37 | download | | Paper | 94.21 | 91.34 | 90.95 | 92.17 | |
Dice Coefficient of 3D-UCaps on hippocampus dataset in 4-folds cross-validation:
| | Anterior | Posterior | Average | Pretrained model | |-------|:--------:|:---------:|:-------:|------------------| | Fold 0 | 86.33 | 83.79 | 85.06 | download | | Fold 1 | 86.57 | 84.51 | 85.54 | download | | Fold 2 | 84.29 | 83.23 | 83.76 | download | | Fold 3 | 85.71 | 83.53 | 84.62 | download | | Mean | 85.73 | 83.77 | 84.75 | | | Paper | 85.07 | 82.49 | 83.78 | |
Result of 3D-UCaps on the cardiac dataset in 4-folds cross-validation:
| | Recall | Precision | Dice | Pretrained model | |-------|:------:|:---------:|:----:|------------------| | Fold 0 | 91.38 | 89.66 | 90.51 | download | | Fold 1 | 89.68 | 95.10 | 91.76 | download | | Fold 2 | 93.12 | 93.00 | 92.53 | download | | Fold 3 | 91.55 | 94.84 | 90.89 | download | | Mean | 91.43 | 93.15 | 91.42 | | | Paper | 92.69 | 89.45 | 90.82 | |
Acknowledgement
The implementation of dynamic routing algorithm and capsule layers were based on the Tensorflow build of CapsNet by its authors in this link
Issues
RuntimeError: Given groups=1, weight of size [16, 2, 5, 5, 5], expected input[1, 1, 32, 32, 32] to have 2 channels, but got 1 channels instead
opened on 2022-10-12 07:47:42 by 217dalaoFile "/home/mtc206/anaconda3/envs/lcj/lib/python3.6/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 219, in validation_step return self.model.validation_step(args, kwargs) File "/home/mtc206/0qsw/SSL4MIS-master/code/3D-UCaps/3D-UCaps-main/module/segcaps.py", line 203, in validation_step overlap=self.overlap, File "/home/mtc206/anaconda3/envs/lcj/lib/python3.6/site-packages/monai/inferers/utils.py", line 130, in sliding_window_inference seg_prob = predictor(window_data, args, kwargs).to(device) # batched patch segmentation File "/home/mtc206/0qsw/SSL4MIS-master/code/3D-UCaps/3D-UCaps-main/module/segcaps.py", line 124, in forward x = self.feature_extractor(x) File "/home/mtc206/anaconda3/envs/lcj/lib/python3.6/site-packages/torch/nn/modules/module.py", line 727, in _call_impl result = self.forward(*input, kwargs) File "/home/mtc206/anaconda3/envs/lcj/lib/python3.6/site-packages/torch/nn/modules/container.py", line 117, in forward input = module(input) File "/home/mtc206/anaconda3/envs/lcj/lib/python3.6/site-packages/torch/nn/modules/module.py", line 727, in _call_impl result = self.forward(input, kwargs) File "/home/mtc206/anaconda3/envs/lcj/lib/python3.6/site-packages/torch/nn/modules/container.py", line 117, in forward input = module(input) File "/home/mtc206/anaconda3/envs/lcj/lib/python3.6/site-packages/torch/nn/modules/module.py", line 727, in _call_impl result = self.forward(input, **kwargs) File "/home/mtc206/anaconda3/envs/lcj/lib/python3.6/site-packages/torch/nn/modules/conv.py", line 573, in forward self.padding, self.dilation, self.groups) RuntimeError: Given groups=1, weight of size [16, 2, 5, 5, 5], expected input[1, 1, 32, 32, 32] to have 2 channels, but got 1 channels instead
when running the code，how can i input the data with 2 channels？
Complexity of the model
opened on 2022-08-28 12:10:06 by cugwuHi, just for reproducibility purposes. In the paper's final model, what are the number of parameters, FLOPs, and inference time? I want to check if I'm recreating a similar situation with your code. Thanks in advance.
RuntimeError: CUDA error: device-side assert triggered
opened on 2022-06-24 06:36:04 by jsong0041First of all, congratulations for your recent paper '3D-UCaps: 3D Capsules Unet for Volumetric Image Segmentation' accepted by MICCAI'21, it's really a great job, and thank you very much for your open source code in github.
As for codes, I used a new dataset as inputs with .tif format, but following errors are thrown:
Validation sanity check: 0%| | 0/1 [00:00<?, ?it/s]C
:/w/b/windows/pytorch/aten/src/ATen/native/cuda/ScatterGatherKernel.cu:312: block: [1926,0,0], thread: [32,0,0] Assertion idx_dim >= 0 && idx_dim < index_size && "index out of bounds" failed.
...
=== Transform input info -- AsDiscrete ===
Traceback (most recent call last):
File "C:\Python36\lib\site-packages\monai\transforms\transform.py", line 84, in apply_transform
return apply_transform(transform, data, unpack_items)
File "C:\Python36\lib\site-packages\monai\transforms\transform.py", line 52, in _apply_transform
return transform(parameters)
File "C:\Python36\lib\site-packages\monai\transforms\post\array.py", line 174, in __call__
img = one_hot(img, num_classes=_nclasses, dim=0)
File "C:\Python36\lib\site-packages\monai\networks\utils.py", line 86, in one_hot
labels = o.scatter(dim=dim, index=labels.long(), value=1)
RuntimeError: CUDA error: device-side assert triggered
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "train.py", line 129, in
Any help is much appreciated.
RuntimeError: CUDA error: an illegal memory access was encountered
opened on 2021-12-08 07:23:06 by wsoniaI deploy the same environment and use the public cardiac data to run the code. But got this problem while training:
Validation sanity check: 0%| | 0/1 [00:00<?, ?it/s]Traceback (most recent call last):
File "train.py", line 137, in
Getting stuck at validation step
opened on 2021-11-18 18:27:07 by cndu234Epoch 0: 0%| | 0/10 [00:00<?, ?it/s]Trying to infer the batch_size from an ambiguous collection. The batch size we found is 1. To avoid any miscalculations, use self.log(..., batch_size=batch_size).
The default behavior for interpolate/upsample with float scale_factor changed in 1.6.0 to align with other frameworks/libraries, and now uses scale_factor directly, instead of relying on the computed output size. If you wish to restore the old behavior, please set recompute_scale_factor=True. See the documentation of nn.Upsample for details.
Epoch 0: 70%|█████████ | 7/10 [00:14<00:06, 2.12s/it, loss=0.724, v_num=11]
Validating: 0it [00:00, ?it/s]
Validating: 0%| | 0/3 [00:00<?, ?it/s]
I have been trying for larger dataset unlike the error here. But im always getting stuck at validation stage. I tried with hippocampus dataset, the results are fine. But with my custom data, Im facing this problem. What could be the reason?
How to compute the metrics between testset predictions and true labels?
opened on 2021-11-18 07:53:55 by cndu234I am using my custom data...After training, how can I compute the metrics between test set predictions and true labels? I am using Hippocampus data loader provided by you . But i have imagesTr, labelsTr for training , and imagesTs, labelsTs for testing. I want to compute metrics for the test set
VinAI Research
GitHub Repository
medical-image-segmentation capsule-network luna16 iseg-challenge hippocampus cardiac-segmentation capsnet segcaps unet