3D UX-Net: A Large Kernel Volumetric ConvNet Modernizing Hierarchical Transformer for Medical Image Segmentation
Official Pytorch implementation of 3D UX-Net, from the following paper:
3D UX-Net: A Large Kernel Volumetric ConvNet Modernizing Hierarchical Transformer for Medical Image Segmentation. Arxiv 2022
Ho Hin Lee, Shunxing Bao, Yuankai Huo, Bennet A. Landman
Vanderbilt University
[arXiv
]
We propose 3D UX-Net, a pure volumetric convolutional network to adapt hierarchical transformers behaviour (e.g. Swin Transformer) for Medical Image Segmentation with less model parameters.
- FeTA 2021, FLARE 2021 Training Code (training.md)
- AMOS 2022 Finetuning Code (finetuning.md)
(Feel free to post suggestions in issues of recommending latest proposed transformer network for comparison. Currently, the network folder is to put the current SOTA transformer. We can further add the recommended network in it for training.)
Methods | resolution | #params | FLOPs | Mean Dice (FeTA2021) | Mean Dice (FLARE2021) |
---|---|---|---|---|---|
TransBTS | 96x96x96 | 31.6M | 110.4G | 0.868 | 0.902 |
UNETR | 96x96x96 | 92.8M | 82.6G | 0.860 | 0.886 |
nnFormer | 96x96x96 | 149.3M | 240.2G | 0.863 | 0.906 |
SwinUNETR | 96x96x96 | 62.2M | 328.4M | 0.867 | 0.929 |
3D UX-Net | 96x96x96 | 53.0M | 639.4G | 0.874 (kernel=7) | 0.938 (kernel=13) |
Methods | resolution | #params | FLOPs | Mean Dice (AMOS2022) |
---|---|---|---|---|
TransBTS | 96x96x96 | 31.6M | 110.4G | 0.792 |
UNETR | 96x96x96 | 92.8M | 82.6G | 0.762 |
nnFormer | 96x96x96 | 149.3M | 240.2G | 0.790 |
SwinUNETR | 96x96x96 | 62.2M | 328.4M | 0.880 |
3D UX-Net | 96x96x96 | 53.0M | 639.4G | 0.900 (kernel=7) |
Training and fine-tuning instructions are in TRAINING.md
Efficient evaulation can be performed for the above three public datasets as follows:
python test_seg.py --root path_to_image_folder --output path_to_output \
--dataset flare --network 3DUXNET --trained_weights path_to_trained_weights \
--mode test --sw_batch_size 4 --overlap 0.7 --gpu 0 --cache_rate 0.2 \
This repository is built using the timm library.
This project is released under the MIT license. Please see the LICENSE file for more information.
If you find this repository helpful, please consider citing:
@Article{lee20223DUX-Net,
author = {Ho Hin Lee and Shunxing Bao and Yuankai Huo and Bennet A. Landman},
title = {3D UX-Net: A Large Kernel Volumetric ConvNet Modernizing Hierarchical Transformer for Medical Image Segmentation]},
journal = {arXiv},
year = {2022},
}