This repo is code of WaveU3S: A Lightweight Wavelet Dual-attention Unet For 3D Medical Image Segmentation (ISBI2024).
WaveU3S is a segmentation network for 3D medical image. Its purpose is to reducing the computation burden while maintaining the performance under 3D segmentation tasks. We use nnUNet as default data preprocessing method and training framework. Comparative experiments are conducted on Flare22 and ACDC datasets. For more details, please refer to our paper.
Our models are built based on nnUNet V2
The model files are located in ./nnunetv2/training/nnUNetTrainer/model
git clone [email protected]:kingofengineer/WaveU3S.git
cd WaveU3S/
pip install -e .You can obtain the datasets via following links：
The pseudo-labels for the 2,000 examples were generated through inference using Hang et al.'s model.
Please refer to nnUNetv2 for data format .
Experiment planning and data preprocessing is to use:
nnUNetv2_plan_and_preprocess -d DATASET_ID --verify_dataset_integrity
Modify the patch size in the configuration file nnUNetPlans.json using the modify_Plans.py .
Training command:
nnUNetv2_train -d DATASET_ID 3d_fullres -f flod WaveU3S
Inference command:
nnUNetv2_predict -i input_dir -o output_dir -d DATASET_ID -c 3d_fullres --save_probabilities -f flod -m WaveU3S
- If you find this work is helpful, please cite our paper
@inproceedings{Z2024,
title={WaveU3S: A Lightweight Wavelet Dual-attention Unet For 3D Medical Image Segmentation},
author={Tianzhao Zhong, Huaishui Yang, Jihao Li, Mengye Lyu, Shaojun Liu},
booktitle={IEEE International Symposium on Biomedical Imaging},