3D TransUNet: Advancing Medical Image Segmentation through Vision Transformers

TL;DR

This paper introduces 3D TransUNet, a novel architecture combining Transformers with U-Net for improved medical image segmentation, effectively capturing global context and enhancing performance on diverse tasks.

Contribution

It extends 2D TransUNet to 3D, integrating Transformer encoder and decoder components into nnU-Net, and demonstrates their effectiveness across multiple medical segmentation tasks.

Findings

Transformer encoder improves multi-organ segmentation accuracy.

Transformer decoder enhances small target and tumor segmentation.

TransUNet outperforms existing methods in various medical imaging benchmarks.

Abstract

Medical image segmentation plays a crucial role in advancing healthcare systems for disease diagnosis and treatment planning. The u-shaped architecture, popularly known as U-Net, has proven highly successful for various medical image segmentation tasks. However, U-Net's convolution-based operations inherently limit its ability to model long-range dependencies effectively. To address these limitations, researchers have turned to Transformers, renowned for their global self-attention mechanisms, as alternative architectures. One popular network is our previous TransUNet, which leverages Transformers' self-attention to complement U-Net's localized information with the global context. In this paper, we extend the 2D TransUNet architecture to a 3D network by building upon the state-of-the-art nnU-Net architecture, and fully exploring Transformers' potential in both the encoder and decoder design. We introduce two key components: 1) A Transformer encoder that tokenizes image patches from a convolution neural network (CNN) feature map, enabling the extraction of global contexts, and 2) A Transformer decoder that adaptively refines candidate regions by utilizing cross-attention between candidate proposals and U-Net features. Our investigations reveal that different medical tasks benefit from distinct architectural designs. The Transformer encoder excels in multi-organ segmentation, where the relationship among organs is crucial. On the other hand, the Transformer decoder proves more beneficial for dealing with small and challenging segmented targets such as tumor segmentation. Extensive experiments showcase the significant potential of integrating a Transformer-based encoder and decoder into the u-shaped medical image segmentation architecture. TransUNet outperforms competitors in various medical applications.