TLUnet: Efficient and Accurate 3D Medical Image Segmentation

Architecture overview of UNETR++

Installation

The code is tested with PyTorch 1.11.0 and CUDA 11.3. After cloning the repository, follow the below steps for installation,

Create and activate conda environment

conda create --name unetr_pp python=3.8
conda activate unetr_pp

Install PyTorch and torchvision

pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 --extra-index-url https://download.pytorch.org/whl/cu113

Install other dependencies

pip install -r requirements.txt

Dataset

We follow the same dataset preprocessing as in UNETR++. We conducted extensive experiments on Synapse

The dataset folders for Synapse should be organized as follows:

./DATASET_Synapse/
  ├── unetr_pp_raw/
      ├── unetr_pp_raw_data/
           ├── Task02_Synapse/
              ├── imagesTr/
              ├── imagesTs/
              ├── labelsTr/
              ├── labelsTs/
              ├── dataset.json
           ├── Task002_Synapse
       ├── unetr_pp_cropped_data/
           ├── Task002_Synapse

The dataset folders for ACDC should be organized as follows:

Please refer to Setting up the datasets on nnFormer repository for more details. Alternatively, you can download the preprocessed dataset for Synapse

Training

The following scripts can be used for training our UNETR++ model on the datasets:

bash training_scripts/run_training_synapse.sh

Evaluation

To reproduce the results of TLUnet:

1- Download Synapse weights and paste model_final_checkpoint.model in the following path:

tlunet/evaluation/unetr_pp_synapse_checkpoint/unetr_pp/3d_fullres/Task002_Synapse/unetr_pp_trainer_synapse__unetr_pp_Plansv2.1/fold_0/

Then, run

bash evaluation_scripts/run_evaluation_synapse.sh

Acknowledgement

This repository is built based on UNETR++ repository.

QuyLe-Minh / TLUnet