Factorizer for Stoke Lesion Segmentation

This repository contains the code and pre-trained models for our method based on Factorizer (repo) submitted to the ISLES'22 challenge for the task of stroke lesion segmentation. The submitted model was an ensemble of Swin Factorizer and Res-U-Net (each described in the paper). This model ranked among top 3 in the ISLES'22 final leaderboard.

Installation

1. Clone the repository:

git clone https://github.com/pashtari/factorizer-isles22.git
cd factorizer-isles22

2. Install dependencies:

pip install -r requirements.txt

Data Preparation

Before training models, we first need to prepare the dataset by taking the following steps:

1. Register and download the official ISLES'22 dataset from this link. The dataset directory have the following structure:

ISLES22
├── dataset_description.json
├── LICENSE
├── participants.tsv
├── README
├── derivatives
│   ├── sub-strokecase0001
│   ├── sub-strokecase0002
│   └── ...
├── sub-strokecase0001
├── sub-strokecase0002
└── ...

2. Perform FLAIR-DWI co-registration:

$ python flair-dwi_registration.py <dataset_dir>

where <dataset_dir> is the dataset directory. The script uses SimpleElastix to align FLAIR images to the DWI space, generating files like {id}_ses-0001_flair_registered.nii.gz.

3. Download datalist.json and place it in the dataset folder:

<dataset_dir>
├── datalist.json # <--
├── dataset_description.json
├── LICENSE
├── participants.tsv
├── README
├── derivatives
│   ├── sub-strokecase0001
│   ├── sub-strokecase0002
│   └── ...
├── sub-strokecase0001
├── sub-strokecase0002
└── ...

Training

To train Swin Factorizer on the first fold of a 5-fold cross-validation on two GPUs, use the following command:

$ python train.py --config configs/config_isles22_fold0_swin-factorizer.yaml

where config_isles22_fold0_swin-factorizer.yaml is a config file. You can find all the config files in ./configs, but before using them, change their values of data_properties to yours. The model checkpoint will then be saved in ./logs/fold0/swin-factorizer/version_0/checkpoints for this example.

Similarly, to train Res-U-Net on the first fold, use the following command:

$ python train.py --config configs/config_isles22_fold0_resunet.yaml

Pre-Trained Models

The Swin Factorizer and Res-U-Net models pre-trained on the ISLES22 dataset are provided in the following. The fold corresponding to each patient is provided in the JSON file.

Model	#Params (M)	FLOPs (G)	Dice (%)	Fold	Config	Checkpoint
Res-U-Net	28.8	145	76.6	0	link	link
Res-U-Net	28.8	145	80.2	1	link	link
Res-U-Net	28.8	145	77.4	2	link	link
Res-U-Net	28.8	145	81.2	3	link	link
Res-U-Net	28.8	145	75.5	4	link	link
Swin Factorizer	7.3	29	76.2	0	link	link
Swin Factorizer	7.3	29	78.9	1	link	link
Swin Factorizer	7.3	29	78.5	2	link	link
Swin Factorizer	7.3	29	80.6	3	link	link
Swin Factorizer	7.3	29	78.6	4	link	link

Inference (CPU)

1. Download all the 10 pre-trained models via the Dropbox link as a folder named logs and place it inside the home directory, that is the structure of folders will be as follows:

factorizer-isles22
├── configs
│   ├── config_isles_fold0_swin-factorizer.yaml
│   └── ...
├── logs
│   ├── fold0
│   │   ├── swin-factorizer
│   │   │   └── version_0
│   │   │       ├── checkpoints
│   │   │       │   └── epoch=515-step=99999.ckpt
│   │   │       └── ...
│   │   └── ...
│   └── ...
├── train.py
├── predict.py
└── ...

2. To make predictions for a test case using the ensemble of all the 10 pre-trained models, run the following command:

$ python predict.py --dwi path/to/dwi --adc path/to/adc --flair path/to/flair --output <output_dir> 

If you want to use only some of the models rather than all of them, specify the argument --checkpoints with a list of paths to the model checkpoints.

Docker

Build Docker Image

1. Clone this repository:

$ git clone https://github.com/pashtari/factorizer-isles22.git

2. Download all the 10 pre-trained models via the Dropbox link as a folder named logs and place it inside the home directory, that is the structure of folders will be as follows:

factorizer-isles22
├── configs
│   ├── config_isles_fold0_swin-factorizer.yaml
│   └── ...
├── logs
│   ├── fold0
│   │   ├── swin-factorizer
│   │   │   └── version_0
│   │   │       ├── checkpoints
│   │   │       │   └── epoch=515-step=99999.ckpt
│   │   │       └── ...
│   │   └── ...
│   └── ...
├── train.py
├── predict.py
└── ...

3. Build docker image:

$ docker build -t factorizer-isles22 ~ 

Run Docker Image

For test images and their predictions organized into the following folder structure:

data
├── input
│   ├── sub-strokecase0001
│   │   ├── sub-strokecase0001_ses-0001_dwi.nii.gz
│   │   ├── sub-strokecase0001_ses-0001_adc.nii.gz
│   │   └── sub-strokecase0001_ses-0001_flair.nii.gz
│   └── ...
└── output
    └── sub-strokecase0001
        ├── sub-strokecase0001_ses-0001_pred.nii.gz
        └── ...

The inference can be performed as follows:

$ cd data
$ docker run -v ${PWD}/input:/input -v ${PWD}/output:/output factorizer-isles22 --dwi /input/sub-strokecase0001/sub-strokecase0001_ses-0001_dwi.nii.gz --adc /input/sub-strokecase0001/sub-strokecase0001_ses-0001_adc.nii.gz --flair /input/sub-strokecase0001/sub-strokecase0001_ses-0001_flair.nii.gz --output /output/sub-strokecase0001/sub-strokecase0001_ses-0001_pred.nii.gz

Save Docker image:

$ docker save factorizer-isles22:latest | gzip > factorizer-isles22.tar.gz

License

This repository is released under the Apache 2.0 license as found in the LICENSE file.

Contact

This repo is currently maintained by Pooya Ashtari (@pashtari).