DOCSEARCHbyAlgolia
Submit Docs

UNETR++ Documentation

repository·main·Indexed 19 days ago

https://github.com/amshaker/unetr_plus_plus

A high-efficiency 3D medical image segmentation framework featuring a hierarchical encoder-decoder structure and a novel Efficient Paired Attention (EPA) block. Designed to reduce the quadratic complexity of standard transformer self-attention, it supports datasets including Synapse, ACDC, Decathlon-Lung, and BRaTs. The framework is tested with PyTorch 1.11.0 and CUDA 11.3.

Tokens
1.8K
Snippets
4
Records
5
Agent score
18%

What's inside UNETR++

  1. What is the UNETR++ architecture?

    main

    UNETR++ is a 3D medical image segmentation model featuring a hierarchical encoder-decoder structure.

    Key components include:

    • Encoder: Processes 3D patches.
    • Decoder: Connected to the encoder via skip connections and convolutional blocks.
    • Efficient Paired-Attention (EPA) Block: The core innovation. It uses parallel attention modules with shared keys-queries and different value layers.
      • The spatial attention module aggregates spatial features with linear complexity relative to input sequence length.
      • The channel attention module emphasizes channel dependencies.
      • The outputs are fused and passed through convolutional blocks to enhance feature representation.

  2. Organize datasets for UNETR++

    main

    UNETR++ follows the preprocessing structure used by nnFormer. Datasets must be organized into specific directory hierarchies containing unetr_pp_raw and unetr_pp_cropped_data folders.

    Each dataset (Synapse, ACDC, Decathlon-Lung, or BRaTs) requires a specific folder name and internal task ID (e.g., Task02_Synapse for Synapse). Within these task folders, you must include imagesTr, imagesTs, labelsTr, labelsTs, and a dataset.json file.

    For detailed preprocessing instructions, refer to the nnFormer dataset setup guide. Alternatively, you can download preprocessed datasets for Synapse, ACDC, Decathlon-Lung, and BRaTs from the provided project links and extract them under the project directory.

    # Example structure for Synapse
./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

  3. Train the UNETR++ model

    main

    Training is performed using provided shell scripts located in the training_scripts/ directory. Use the script corresponding to your specific dataset:

    • Synapse: bash training_scripts/run_training_synapse.sh
    • ACDC: bash training_scripts/run_training_acdc.sh
    • Decathlon-Lung: bash training_scripts/run_training_lung.sh
    • BRaTs (Tumor): bash training_scripts/run_training_tumor.sh
    bash training_scripts/run_training_synapse.sh
bash training_scripts/run_training_acdc.sh
bash training_scripts/run_training_lung.sh
bash training_scripts/run_training_tumor.sh

  4. Install UNETR++

    main

    To install UNETR++, create a conda environment and install the required PyTorch and dependency versions. The code is tested with PyTorch 1.11.0 and CUDA 11.3.

    1. Create and activate a conda environment named unetr_pp with Python 3.8.
    2. Install torch==1.11.0+cu113 and torchvision==0.12.0+cu113 using the specific CUDA 11.3 index URL.
    3. Install remaining dependencies via requirements.txt.
    # 1. Create and activate conda environment
conda create --name unetr_pp python=3.8
conda activate unetr_pp

# 2. 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

# 3. Install other dependencies
pip install -r requirements.txt

  5. Reproduce UNETR++ results via evaluation scripts

    main

    To reproduce the reported results for UNETR++, you must download the pre-trained model weights for specific datasets, place the model_final_checkpoint.model file in the correct directory structure, and execute the corresponding evaluation shell script.

    Supported datasets include:

    • Synapse
    • ACDC
    • Decathlon-Lung
    • BRaTs (Tumor segmentation)
    ### Synapse
1. Download [Synapse weights](https://drive.google.com/file/d/13JuLMeDQRR_a3c3tr2V2oav6I29fJoBa)
2. Place `model_final_checkpoint.model` in:
`unetr_pp/evaluation/unetr_pp_synapse_checkpoint/unetr_pp/3d_fullres/Task002_Synapse/unetr_pp_trainer_synapse__unetr_pp_Plansv2.1/fold_0/`
3. Run:
```bash
bash evaluation_scripts/run_evaluation_synapse.sh

    ACDC

    1. Download ACDC weights
    2. Place model_final_checkpoint.model in:

    unetr_pp/evaluation/unetr_pp_acdc_checkpoint/unetr_pp/3d_fullres/Task001_ACDC/unetr_pp_trainer_acdc__unetr_pp_Plansv2.1/fold_0/ 3. Run:

    bash evaluation_scripts/run_evaluation_acdc.sh

    Decathlon-Lung

    1. Download Decathlon-Lung weights
    2. Place model_final_checkpoint.model in:

    unetr_pp/evaluation/unetr_pp_lung_checkpoint/unetr_pp/3d_fullres/Task006_Lung/unetr_pp_trainer_lung__unetr_pp_Plansv2.1/fold_0/ 3. Run:

    bash evaluation_scripts/run_evaluation_lung.sh

    BRaTs

    1. Download BRaTs weights
    2. Place model_final_checkpoint.model in:

    unetr_pp/evaluation/unetr_pp_lung_checkpoint/unetr_pp/3d_fullres/Task003_tumor/unetr_pp_trainer_tumor__unetr_pp_Plansv2.1/fold_0/ 3. Run:

    bash evaluation_scripts/run_evaluation_tumor.sh