Official PyTorch implementation of the paper RWKV-UNet: Improving UNet with Long-Range Cooperation for Effective Medical Image Segmentation.
Synapse (BTCV preprocessed data) and ACDC data are available at TransUNet's repo.
The directory structure of the whole project is as follows:
.
├── RWKVUNet
│ ├──datasets
│ │ └── dataset_*.py
│ ├──train.py
│ ├──test.py
│ ├──...
│ └──data
│ └── Synapse
│ │ ├── test_vol_h5
│ │ │ ├── case0001.npy.h5
│ │ │ └── *.npy.h5
│ │ └── train_npz
│ │ ├── case0005_slice000.npz
│ │ └── *.npz
│ │
│ └──ACDC
│ ├── ACDC_training_volumes
│ │ ├── patient100_frame01.h5
│ │ └── *.h5
│ └── ACDC_training_slices
│ ├── patient100_frame13_slice_0.h5
│ └── *.h5
Please prepare an environment with python=3.9, and then use the command "pip install -r requirements.txt" for the dependencies.
Pretrained weights for the encoder can be downloaded at (https://drive.google.com/drive/folders/1odF_NK5wYRkE0C3w9eoLUQEVbxefj66e?usp=sharing).
Checkpoints for RWKV-UNet can be downloaded at (https://drive.google.com/drive/folders/19y_8Mzmw5u6Bg-iVfmh6-vRCBDdy149_?usp=sharing).
- Run the training script on the Synapse dataset.
python train.py --dataset Synapse --max_epochs 30 --base_lr 0.001 --img_size 224 --pretrained_path netB.pth- Run the test script on the Synapse dataset.
python test.py --dataset Synapse --max_epochs 30 --base_lr 0.001 --img_size 224- You can also specify the weights for inference by yourself.
python test.py --dataset Synapse --path_specific ckpts/Synapse_base.pthThis code base uses certain code blocks and helper functions from TransUNet and Vision-RWKV.
@misc{jiang2025rwkvunetimprovingunetlongrange,
title={RWKV-UNet: Improving UNet with Long-Range Cooperation for Effective Medical Image Segmentation},
author={Juntao Jiang and Jiangning Zhang and Weixuan Liu and Muxuan Gao and Xiaobin Hu and Zhucun Xue and Yong Liu and Shuicheng Yan},
year={2025},
eprint={2501.08458},
archivePrefix={arXiv},
primaryClass={eess.IV},
url={https://arxiv.org/abs/2501.08458},
}