SeUNet-Trans: A Simple yet Effective UNet-Transformer Model for Medical Image Segmentation

TL;DR

SeUNet-Trans combines UNet and Transformer architectures to improve medical image segmentation by capturing both local and global features efficiently, demonstrating superior performance across multiple datasets.

Contribution

The paper introduces a simple yet effective UNet-Transformer model that integrates pixel-level embedding and spatial-reduction attention for enhanced medical image segmentation.

Findings

Outperforms state-of-the-art models on seven datasets

Effectively captures long-range dependencies in images

Reduces computational overhead with spatial-reduction attention

Abstract

Automated medical image segmentation is becoming increasingly crucial to modern clinical practice, driven by the growing demand for precise diagnosis, the push towards personalized treatment plans, and the advancements in machine learning algorithms, especially the incorporation of deep learning methods. While convolutional neural networks (CNN) have been prevalent among these methods, the remarkable potential of Transformer-based models for computer vision tasks is gaining more acknowledgment. To harness the advantages of both CNN-based and Transformer-based models, we propose a simple yet effective UNet-Transformer (seUNet-Trans) model for medical image segmentation. In our approach, the UNet model is designed as a feature extractor to generate multiple feature maps from the input images, then the maps are propagated into a bridge layer, which is introduced to sequentially connect the UNet and the Transformer. In this stage, we approach the pixel-level embedding technique without position embedding vectors, aiming to make the model more efficient. Moreover, we apply spatial-reduction attention in the Transformer to reduce the computational/memory overhead. By leveraging the UNet architecture and the self-attention mechanism, our model not only retains the preservation of both local and global context information but also is capable of capturing long-range dependencies between input elements. The proposed model is extensively experimented on seven medical image segmentation datasets including polyp segmentation to demonstrate its efficacy. Comparison with several state-of-the-art segmentation models on these datasets shows the superior performance of our proposed seUNet-Trans network.