HiFormer - LITS: Hierarchical Multi-scale Representations Using Transformers for Medical Image Segmentation, implemented on the LITS Dataset
This is a fork of the original HiFormer repository which I have modified to implement the Liver Tumor Segmentation (LITS) dataset, instead of the original Synapse Dataset.
All model credits and original implementation credits go to the original authors of the paper.
This code has been implemented in Python 3 using PyTorch, and tested on Ubuntu OS, and popular online data science platforms Kaggle and Google Colab.
The train.py and test.py scripts require a root path directory. The structure of this directory should be as follows:
root_path
├── liver_0
│ ├── images
│ │ └── 45.jpg
│ │ └── 46.jpg
| | └── ---
│ └── masks
│ ├── liver
| │ └── 46.jpg
| │ └── 47.jpg
| | └── ---
│ └── cancer
| └── 46.jpg
| └── 47.jpg
| └── ---
|
|
├── liver_130
│ ├── images
│ │ └── 45.jpg
│ │ └── 46.jpg
| | └── ---
│ └── masks
│ ├── liver
| │ └── 46.jpg
| │ └── 47.jpg
| | └── ---
│ └── cancer
| └── 46.jpg
| └── 47.jpg
| └── ---
Note that separate folders for train, test, and split are not required!
Based on the random seed given to the train and test script, it will automatically divide the 131 folders into train and split on a 0.8,0.1,0.1 basis.
MAKE SURE TO GIVE THE SAME SEED TO TRAIN AND TEST SCRIPT, else the testing split for test.py may include a folder present in the training split of train.py.
Please follow the instructions below to run the script -
-
Obtain the dataset in jpg format and create the root structure as defined above
-
Run the following code to install the Requirements.
pip install -r requirements.txt -
Run the below code to train HiFormer on the synapse dataset.
python train.py --root_path /LITS_Data/ --model_name hiformer-b --batch_size 16 --eval_interval 10 --max_epochs 400 --seed 69
--root_path [Root data path]
--eval_interval [Evaluation epoch]
--model_name [Choose from [hiformer-s, hiformer-b, hiformer-l]]
--is_liver [If added, model will train/test using the liver images. If not added, it will train/test on the cancer(tumor) images]
--seed [Numpy and Torch random seed for reproducibility, and to ensure consistent train-test-val split amongst training and testing]
-
Run the below code to test HiFormer on the synapse dataset.
python test.py --test_path /LITS_Data/ --model_name hiformer-b --model_weight './hiformer-b_best.pth' --seed 69--root_path [Root data path]
--model_name [Choose from [hiformer-s, hiformer-b, hiformer-l]]
--is_liver [If added, model will train/test using the liver images. If not added, it will train/test on the cancer(tumor) images]
--seed [Numpy and Torch random seed for reproducibility, and to ensure consistent train-test-val split amongst training and testing]
--model_weight [Path to model weights (.pth) file for testing]
The LITS Implementation has been done by Jasmer Singh Sanjotra.
Original implementations are done by Amirhossein Kazerouni, Milad Soltany and Moein Heidari.
Original Authors -
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.