SLab Medical Segmentation (SMS)

A advanced and general medical segmentation toolkit developed by the SLab team, led by Prof. Shuang Song and Kangneng Zhou.

🚀 News

• (2025.11.20): What a beautiful day! The first version of the SLab Medical Segmentation (SMS) repository has been released.

Table of Contents

Follow steps 1-3 to run our toolkit.

1. Introduction
2. Installation
3. Prepare the Dataset
4. Run the Training
5. Run the Inference
6. Visualization
7. Results
8. To Do List
9. By The Way
10. Acknowledgements

Introduction

Support 3D Networks:

• Unetr (WACV 2022, code)

• unet (MICCAI 2016, code)

• unetr_pp (TMI 2024, code)

• Segformer3D (CVPR-W 2024, code)

• SegMamba (ICCV 2025, code)

• Slim-UNETR (TMI 2023, code)

• vtunet (MICCAI 2022, code)

• AttentionUnet (MIDL 2018, code)

• SwinUNETR (MICCAI-W 2021, code)

• UNet++ (TMI 2019, code)

• 3DUXNET (ICLR 2023, code)

• nnFormer (TIP 2023, code)

• EfficientMedNext (MICCAI 2025, code)

• nnUnet (Nature Methods 2021, code)

• VNet (3DV 2016, code)

• MedFormer (arXiv 2022, code)

• TransBTS (MICCAI2021, code)

Support 2D Networks:

• UNet (MICCAI 2015, code)
• UNet++ (TMI 2019, code)
• DeepLabV3 (arXiv 2017, code)
• FCN (CVPR 2015, code)
• SegNet (TPAMI 2017, code)
• PSPNet (CVPR 2017)
• HighResNet (IPMI 2017, code)
• MiniSeg (TNNLS 2022, code)
• AttentionUnet (MIDL 2018, code)
• DANet (CVPR 2019, code)
• MedFormer (arXiv 2022, code)
• TransUnet (MIA 2024, code)
• SwinUNet (ECCV-W 2024, code)

Support Datasets:

• Brain Tumor Segmentation Dataset (BraTS) - 3D
• DSA (Digital Subtraction Angiography) - 2D

Installation

conda create -n slab_med_seg python=3.10
conda activate slab_med_seg
pip install torch torchvision
pip install -r requirements.txt

Prepare the Dataset

design your dataloader in dataset

Run the Training

For 3D Models:

python train.py --dataset_name your_dataset_name \
                --dataset_path your_dataset_path \
                --batch 4 \
                --num_workers 4 \
                --queue_length 5 \
                --samples_per_volume 5 \
                --patch_size 96 96 96 \
                --expname your_exp_name \
                --model_name your_model_name \
                --in_channels 1 \
                --out_channels 3 \
                --loss_type combine

For 2D Models:

python train.py --dataset_name dsa \
                --dataset_path data/data_2d/DSA \
                --model_name unet \
                --dimension 2d \
                --in_channels 1 \
                --out_channels 1 \
                --batch 4 \
                --num_epochs 100 \
                --expname dsa_unet_2d \
                --loss_type combine

Run the Inference

python inference.py \
    --checkpoint_path logs/exp/best_model.pth \
    --input_path data/test/img/ \
    --output_dir results/ \
    --config logs/exp/config.yaml \
    --use_sliding_window \
    --calculate_metrics \
    --save_probability

Results

We benchmark SLab Medical Segmentation (SMS) both qualitatively and quantitatively against the current state-of-the-art models such as MedNeXt and several other models on three widely used medical datasets: BRaTs2017.

Brain Tumor Segmentation Dataset (BraTs)

Stay tuned!

TODO

• More datasets

By The Way

If you're using this project and have feedback, feel free to contact Kangneng Zhou via WeChat (kangkangellis666) or email (elliszkn@163.com).

Acknowledgements

This repository offers advanced and general medical segmentation in 3D and 2D. The project was completed under the supervision of Prof. Shuang Song and Kangneng Zhou. Special thanks to Haoyu Yuan, Dian Song, Jingdan Zhang, Yihang Xu and Sibo Zhao for their invaluable support.