OrgUNETR: Utilizing Organ Information and Squeeze and Excitation Block for Improved Tumor Segmentation

Abstract

Convolutional Neural Networks (CNNs) have demonstrated remarkable performance in medical image segmentation tasks, with the U-Net architecture being a prominent example. The UNet Transformer (UNETR), an advanced variant of U-Net, incorporates a transformer architecture to effectively capture long-range dependencies in Computed Tomography (CT) scans. However, the application of deep learning models for tumor segmentation remains limited due to the challenges posed by the small size and unpredictable locations of tumors. To address this issue, we propose a novel approach that leverages organ information to improve tumor localization. Our model, named OrgUNETR, incorporates organ context by utilizing the fact that tumors typically exist within specific organs. By integrating organ information, OrgUNETR successfully detects tumors in CT scans with enhanced accuracy. Experimental results demonstrate that OrgUNETR surpasses the performance of a baseline model by achieving a 40.86% improvement in Dice score on the KiTS19 dataset and a 32.69% improvement on the Prostate158 dataset. Furthermore, we optimize the computational efficiency of UNETR by replacing the Multi-Head Self-Attention (MHSA) layers with Squeeze and Excitation (SE) layers, which perform attention in a similar manner. This modification reduces the computational cost by 13.9% while maintaining comparable performance. The proposed OrgUNETR model represents a significant advancement in tumor segmentation, combining the benefits of organ context and efficient attention mechanisms to achieve promising results. This research has the potential to assist medical professionals in accurate tumor localization and improve patient outcomes in clinical settings.

Keywords

• Organ segmentation
• tumor segmentation
• medical segmentation
• deep learning
• squeeze and excitation network
• transformer