Medical-Transformer
Pytorch Code for the paper "Medical Transformer: Gated Axial-Attention for Medical Image Segmentation"
About this repo:
This repo hosts the code for the following networks:
- Gated Axial Attention U-Net
- MedT
Introduction
Majority of existing Transformer-based network architectures proposed for vision applications require large-scale datasets to train properly. However, compared to the datasets for vision applications, for medical imaging the number of data samples is relatively low, making it difficult to efficiently train transformers for medical appli- cations. To this end, we propose a Gated Axial-Attention model which extends the existing architectures by introducing an additional control mechanism in the self-attention module. Furthermore, to train the model effectively on medical images, we propose a Local-Global training strat- egy (LoGo) which further improves the performance. Specifically, we op- erate on the whole image and patches to learn global and local features, respectively. The proposed Medical Transformer (MedT) uses LoGo training strategy on Gated Axial Attention U-Net.
Using the code:
- Clone this repository:
git clone https://github.com/jeya-maria-jose/Medical-Transformer
cd Medical-Transformer
The code is stable using Python 3.6.10, Pytorch 1.4.0
To install all the dependencies using conda:
conda env create -f environment.yml
conda activate medt
To install all the dependencies using pip:
pip install -r requirements.txt
Links for downloading the public Datasets:
- GLAS Dataset - Link (Original) | Link (Resized)
- MoNuSeG Dataset - Link (Original)
- Brain Anatomy US dataset from the paper will be made public soon !
Using the Code for your dataset
Dataset Preparation
Prepare the dataset in the following format for easy use of the code. The train and test folders should contain two subfolders each: img and label. Make sure the images their corresponding segmentation masks are placed under these folders and have the same name for easy correspondance. Please change the data loaders to your need if you prefer not preparing the dataset in this format.
Train Folder-----
img----
0001.png
0002.png
.......
label---
0001.png
0002.png
.......
Validation Folder-----
img----
0001.png
0002.png
.......
label---
0001.png
0002.png
.......
Test Folder-----
img----
0001.png
0002.png
.......
label---
0001.png
0002.png
.......
- The ground truth images should have pixels corresponding to the labels. Example: In case of binary segmentation, the pixels in the GT should be 0 or 255.
Training Command:
python train.py --train_dataset "enter train directory" --val_dataset "enter validation directory" --direc 'path for results to be saved' --batch_size 4 --epoch 400 --save_freq 10 --modelname "gatedaxialunet" --learning_rate 0.001 --imgsize 128 --gray "no"
Change modelname to MedT or logo to train them
Testing Command:
python test.py --loaddirec "./saved_model_path/model_name.pth" --val_dataset "test dataset directory" --direc 'path for results to be saved' --batch_size 1 --modelname "gatedaxialunet" --imgsize 128 --gray "no"
The results including predicted segmentations maps will be placed in the results folder along with the model weights. Run the performance metrics code in MATLAB for calculating F1 Score and mIoU.
Notes:
Note that these experiments were conducted in Nvidia Quadro 8000 with 48 GB memory.
Acknowledgement:
The dataloader code is inspired from pytorch-UNet . The axial attention code is developed from axial-deeplab.
Citation:
To add
Open an issue or mail me directly in case of any queries or suggestions.
Thank you.
360 Dec 10, 2022
94 Jan 6, 2023
48 Sep 27, 2022
90 Dec 8, 2022
520 Dec 30, 2022
6.9k Jan 3, 2023
119 Nov 24, 2022
1.4k Jan 1, 2023
9.2k Jan 2, 2023
121 Dec 17, 2022
172 Dec 12, 2022
111 Dec 8, 2022
0 Jul 13, 2021
8 Nov 21, 2022
157 Dec 11, 2022
6 Mar 17, 2022
5 Sep 29, 2021
76 Jan 7, 2023