/HiFormer

HiFormer: Hierarchical Multi-scale Representations Using Transformers for Medical Image Segmentation (WACV 2023)

Primary LanguageJupyter NotebookMIT LicenseMIT

HiFormer: Hierarchical Multi-scale Representations Using Transformers for Medical Image Segmentation, WACV 2023

Hierarchical Multi-scale Representations for medical image segmentation with state of the art results on synapse multi-organ segmentation, skin lesion segmentation, multiple myeloma cell segmentation. This method designs two multi-scale feature representations using the seminal Swin Transformer module and a CNN-based encoder, then using a Double-Level Fusion (DLF) module, it allows a fine fusion of global and local features. If this code helps with your research please consider citing the following paper:

Moein Heidari*, Amirhossein Kazerouni*, Milad Soltany*, R. Azad, Ehsan Khodapanah Aghdam, Julien Cohen-Adad and Dorit Merhof , "HiFormer: Hierarchical Multi-scale Representations Using Transformers for Medical Image Segmentation", download link.

Please consider starring us, if you found it useful. Thanks!

Updates

This code has been implemented in python language using Pytorch library and tested in ubuntu OS, though should be compatible with related environment. following Environement and Library needed to run the code:

  • Python 3
  • Pytorch

Prepare data, Train/Test

Please go to "Instructions.ipynb" for complete detail on dataset preparation and Train/Test procedure or follow the instructions below. Open In Colab

  1. Download the Synapse Dataset from here.

  2. Run the following code to install the Requirements.

    pip install -r requirements.txt

  3. Run the below code to train HiFormer on the synapse dataset.

    python train.py --root_path ./data/Synapse/train_npz --test_path ./data/Synapse/test_vol_h5  --model_name hiformer-b --batch_size 10 --eval_interval 20 --max_epochs 400

    --root_path [Train data path]

    --test_path [Test data path]

    --eval_interval [Evaluation epoch]

    --model_name [Choose from [hiformer-s, hiformer-b, hiformer-l]]

  4. Run the below code to test HiFormer on the synapse dataset.

    python test.py --test_path ./data/Synapse/test_vol_h5 --model_name hiformer-b --is_savenii --model_weight './hiformer-b_best.pth'

    --test_path [Test data path]

    --model_name [choose from [hiformer-s, hiformer-b, hiformer-l]]

    --is_savenii [Whether to save results during inference]

    --model_weight [HiFormer trained model path]

Quick Overview

Diagram of the proposed method

Results

For evaluating the performance of the proposed method, three challenging tasks in medical image segmentaion has been considered. In bellow, results of the proposed approach for synapse multi-organ segmentation is illustrated.

Performance Comparision on Synapse Multi-Organ Segmentation

Methods

DSC↑

HD↓

Aorta

Gallbladder

Kidney(L)

Kidney(R)

Liver

Pancreas

Spleen

Stomach
DARR 69.77 - 74.74 53.77 72.31 73.24 94.08 54.18 89.90 45.96
R50 U-Net 74.68 36.87 87.74 63.66 80.60 78.19 93.74 56.90 85.87 74.16
U-Net 76.85 39.70 89.07 69.72 77.77 68.60 93.43 53.98 86.67 75.58
R50 Att-UNet 75.57 36.97 55.92 63.91 79.20 72.71 93.56 49.37 87.19 74.95
Att-UNet 77.77 36.02 89.55 68.88 77.98 71.11 93.57 58.04 87.30 75.75
R50 ViT 71.29 32.87 73.73 55.13 75.80 72.20 91.51 45.99 81.99 73.95
TransUnet 77.48 31.69 87.23 63.13 81.87 77.02 94.08 55.86 85.08 75.62
SwinUnet 79.13 21.55 85.47 66.53 83.28 79.61 94.29 56.58 90.66 76.60
LeVit-Unet 78.53 16.84 78.53 62.23 84.61 80.25 93.11 59.07 88.86 72.76
DeepLabv3+ (CNN) 77.63 39.95 88.04 66.51 82.76 74.21 91.23 58.32 87.43 73.53
HiFormer-S 80.29 18.85 85.63 73.29 82.39 64.84 94.22 60.84 91.03 78.07
HiFormer-B 80.39 14.70 86.21 65.69 85.23 79.77 94.61 59.52 90.99 81.08
HiFormer-L 80.69 19.14 87.03 68.61 84.23 78.37 94.07 60.77 90.44 82.03

Perceptual visualization results on test data

Synapse Multi-Organ Segmentation result

References

Query

All implementations are done by Amirhossein Kazerouni, Milad Soltany and Moein Heidari. For any query please contact us for more information.

amirhossein477@gmail.com

soltany.m.99@gmail.com

moeinheidari7829@gmail.com

Citation

@inproceedings{heidari2023hiformer,
  title={Hiformer: Hierarchical multi-scale representations using transformers for medical image segmentation},
  author={Heidari, Moein and Kazerouni, Amirhossein and Soltany, Milad and Azad, Reza and Aghdam, Ehsan Khodapanah and Cohen-Adad, Julien and Merhof, Dorit},
  booktitle={Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision},
  pages={6202--6212},
  year={2023}
}