SegMamba

Now we have open-sourced the pre-processing, training, inference, and metrics computation codes.

SegMamba: Long-range Sequential Modeling Mamba For 3D Medical Image Segmentation

Our advantage in speed and memory.

Contact

If you have any questions about our project, please feel free to contact us by email at zxing565@connect.hkust-gz.edu.cn or via WeChat at 18340097191.

Environment install

Clone this repository and navigate to the root directory of the project.

git clone https://github.com/ge-xing/SegMamba.git

cd SegMamba

Install causal-conv1d

cd causal-conv1d

python setup.py install

Install mamba

cd mamba

python setup.py install

Install monai

pip install monai

Simple test

python 0_inference.py

Preprocessing, training, testing, inference, and metrics computation

Data downloading

Data is from https://arxiv.org/abs/2305.17033

Download from Baidu Disk https://pan.baidu.com/s/1C0FUHdDtWNaYWLtDDP9TnA?pwd=ty22提取码ty22

Download from OneDrive https://hkustgz-my.sharepoint.com/:f:/g/personal/zxing565_connect_hkust-gz_edu_cn/EqqaINbHRxREuIj0XGicY2EBv8hjwEFKgFOhF_Ub0mvENw?e=yTpE9B

Preprocessing

In my setting, the data directory of BraTS2023 is : "./data/raw_data/BraTS2023/ASNR-MICCAI-BraTS2023-GLI-Challenge-TrainingData/"

First, we need to run the rename process.

python 1_rename_mri_data.py

Then, we need to run the pre-processing code to do resample, normalization, and crop processes.

python 2_preprocessing_mri.py

After pre-processing, the data structure will be in this format:

Training

When the pre-processing process is done, we can train our model.

We mainly use the pre-processde data from last step: data_dir = "./data/fullres/train"

python 3_train.py

The training logs and checkpoints are saved in: logdir = f"./logs/segmamba"

Inference

When we have trained our models, we can inference all the data in testing set.

python 4_predict.py

When this process is done, the prediction cases will be put in this path: save_path = "./prediction_results/segmamba"

Metrics computation

We can obtain the Dice score and HD95 on each segmentation target (WT, TC, ET for BraTS2023 dataset) using this code:

python 5_compute_metrics.py --pred_name="segmamba"

Acknowledgement

Many thanks for these repos for their great contribution!

https://github.com/MIC-DKFZ/nnUNet

https://github.com/Project-MONAI/MONAI

https://github.com/hustvl/Vim

https://github.com/bowang-lab/U-Mamba