UNETR++ offers an efficient and accurate solution for 3D medical image segmentation, targeting researchers and practitioners in medical imaging. It addresses the computational bottleneck of standard transformer self-attention in volumetric data by introducing a novel Efficient Paired Attention (EPA) block, achieving state-of-the-art results with significantly reduced computational cost.

How It Works

UNETR++ employs a hierarchical encoder-decoder architecture. The core innovation is the Efficient Paired Attention (EPA) block, which uses two inter-dependent attention branches: one with linear complexity for spatial attention and another for channel attention. These branches share weights for query and key mapping, enabling complementary feature learning while reducing parameters. The outputs are fused and processed by convolutional blocks to enhance feature representation for improved segmentation masks.

Quick Start & Requirements

• Install: Create a conda environment (conda create --name unetr_pp python=3.8), activate it, and install PyTorch 1.11.0+cu113 and torchvision 0.12.0+cu113, then pip install -r requirements.txt.
• Prerequisites: PyTorch 1.11.0, CUDA 11.3, Python 3.8.
• Datasets: Requires specific folder organization for Synapse, ACDC, Decathlon-Lung, and BRaTs datasets, or download preprocessed versions.
• Links: nnFormer repository for dataset setup guidance.

Highlighted Details

• Achieves a Dice Score of 87.2% on the Synapse dataset, a new state-of-the-art.
• Demonstrates over 71% reduction in parameters and FLOPs compared to prior best methods.
• Evaluated on five benchmarks: Synapse, BTCV, ACDC, BRaTs, and Decathlon-Lung.
• Qualitative comparisons show improved segmentation accuracy over UNETR, Swin UNETR, and nnFormer.

Maintenance & Community

• Accepted to IEEE TMI-2024.
• Code and weights released for multiple datasets.
• Contact via GitHub issues or email: abdelrahman.youssief@mbzuai.ac.ae.

Licensing & Compatibility

• The repository is built based on the nnFormer repository. The specific license for UNETR++ itself is not explicitly stated in the README, but nnFormer is Apache 2.0 licensed. Compatibility for commercial use should be verified.

Limitations & Caveats

The installation instructions specify PyTorch 1.11.0 and CUDA 11.3, which are older versions. The README does not explicitly state the license for UNETR++ itself, requiring further investigation for commercial use.