Rolling-Unet

The official implementation of the paper "Rolling-Unet: Revitalizing MLP’s Ability to Efficiently Extract Long-Distance Dependencies for Medical Image Segmentation" published in AAAI-2024.

Introduction

Medical image segmentation methods based on deep learning network are mainly divided into CNN and Transformer. However, CNN struggles to capture long-distance dependencies, while Transformer suffers from high computational complexity and poor local feature learning. To efficiently extract and fuse local features and long-range dependencies, this paper proposes Rolling-Unet, which is a CNN model combined with MLP. Specifically, we propose the core R-MLP module, which is responsible for learning the long-distance dependency in a single direction of the whole image. By controlling and combining R-MLP modules in different directions, OR-MLP and DOR-MLP modules are formed to capture long-distance dependencies in multiple directions. Further, Lo2 block is proposed to encode both local context information and long-distance dependencies without excessive computational burden. Lo2 block has the same parameter size and computational complexity as a 3×3 convolution. The experimental results on four public datasets show that Rolling-Unet achieves superior performance compared to the state-of-the-art methods.

Implementation

• The experimental environment of this paper includes: Ubuntu 20.04 LTS system, an i9-12900K CPU, an NVIDIA RTX A6000 GPU, Python version 3.8, Pytorch version 1.13, and CUDA version 11.6.

• Clone this repository:

git clone https://github.com/Jiaoyang45/Rolling-Unet
cd Rolling-Unet

Datasets

1. ISIC 2018 - Link
2. BUSI - Link
3. GLAS - Link
4. CHASEDB1 - Link

Data Format

• Make sure to put the files as the following structure. For binary segmentation, just use folder 0.

inputs
└── <dataset name>
    ├── images
    |   ├── 001.png
    │   ├── 002.png
    │   ├── 003.png
    │   ├── ...
    |
    └── masks
        └── 0
            ├── 001.png
            ├── 002.png
            ├── 003.png
            ├── ...

Training and Validation

• Train the model.

python train.py

• Evaluate.

python val.py

• This code includes visualization of the training process through tensorboardX, and calculation multiple indicators through MedPy.

Acknowledgements

This code repository is implemented based on UNeXt.

References

1. UNet, UNet++, Attention-UNet - link
2. Medical-Transformer - link
3. UCTransNet - link
4. UNeXt - link
5. DconnNet - link

Citations

If this code is helpful for your study, please cite:

1. Liu, Y., Zhu, H., Liu, M., Yu, H., Chen, Z., & Gao, J. (2024). Rolling-Unet: Revitalizing MLP’s Ability to Efficiently Extract Long-Distance Dependencies for Medical Image Segmentation. Proceedings of the AAAI Conference on Artificial Intelligence, 38(4), 3819-3827. https://doi.org/10.1609/aaai.v38i4.28173