SHFormer_pytorch

This repo holds code for An Efficient and Rapid Medical Image Segmentation Network

Setup

Environment

Run the following code to prepare python environment

sh prepare_env.sh

Prepare Data for ISIC-2018

1. Download download the officially provided training set: ISIC-2018
2. run python data_preprocessing/data_tnpz_18.py to process datas

For ISIC-2018 dataset, we resize all images to 512 $\times$ 512 and save them as npz file. After doing this, your dataset directory ISIC2018 should be like this:

ISIC2018
|-- ISIC2018_Task1-2_Training_Input
|-- ISIC2018_Task1_Training_GroundTruth
|-- train
|   |-- ISIC_0000000.npz
|   |-- ...
|   |-- ori_size.pkl
|   |-- splits_all.pkl

The file splits_all.pkl saves the partition of the dataset for ISIC-2018.

Training & Test

Training

usage: python training/run_training.py [--config CONFIG_FILE_NAME] [-gpu GPU_NUM]
Train model.

Arguments:      
  -c CONFIG_FILE_NAME              Configuration file name. (./configs/***.yaml) 
  -g GPU_NUM                       GPU ID(number). Only support single gpu setting.

Note that you need to set the argument validation_only to False in config file.

Test

Similar to Training. Set validation_only to True to enable testing mode. Run the script test.sh to test SHFormer on ISIC-2018 dataset.

Example for load pretrained chechpoint to test model:

python training/run_training.py --config configs/shformer_add.yaml \
    --pretrain_pth pretrained_params/shformer_add.pth

Other Visualization

The visualization code used in the paper is located in the training/trainer/SHFormerTrainer.py section.

Reference

• nnUNet
• GradCAM
• mmsegmentation