Unlocking Fine-Grained Details with Wavelet-based High-Frequency Enhancement in Transformers

arXiv (Cornell University)

Abstract

Medical image segmentation is a critical task that plays a vital role in diagnosis, treatment planning, and disease monitoring. Accurate segmentation of anatomical structures and abnormalities from medical images can aid in the early detection and treatment of various diseases. In this paper, we address the local feature deficiency of the Transformer model by carefully re-designing the self-attention map to produce accurate dense prediction in medical images. To this end, we first apply the wavelet transformation to decompose the input feature map into low-frequency (LF) and high-frequency (HF) subbands. The LF segment is associated with coarse-grained features while the HF components preserve fine-grained features such as texture and edge information. Next, we reformulate the self-attention operation using the efficient Transformer to perform both spatial and context attention on top of the frequency representation. Furthermore, to intensify the importance of the boundary information, we impose an additional attention map by creating a Gaussian pyramid on top of the HF components. Moreover, we propose a multi-scale context enhancement block within skip connections to adaptively model inter-scale dependencies to overcome the semantic gap among stages of the encoder and decoder modules. Throughout comprehensive experiments, we demonstrate the effectiveness of our strategy on multi-organ and skin lesion segmentation benchmarks. The implementation code will be available upon acceptance. \href{https://github.com/mindflow-institue/WaveFormer}{GitHub}.

Authors 7

1. Reza Azad Aachen Faculty of Electrical Engineering and Information Technology

   RWTH Aachen University

   Affiliation as printed

   Faculty of Electrical Engineering and Information Technology , RWTH Aachen University , Germany

2. Amirhossein Kazerouni

   Iran University of Science and Technology

   Affiliation as printed

   School of Electrical Engineering , Iran University of Science and Technology , Iran

3. Alaa Sulaiman

   National University of Malaysia

   Affiliation as printed

   Faculty of Information Science and Technology , Universiti Kebangsaan , Malaysia

4. Afshin Bozorgpour

   University of Regensburg

   Affiliation as printed

   Faculty of Informatics and Data Science , University of Regensburg , Germany

5. Ehsan Khodapanah Aghdam

   Shahid Beheshti University

   Affiliation as printed

   Department of Electrical Engineering , Shahid Beheshti University , Iran

6. Abin Jose Aachen Faculty of Electrical Engineering and Information Technology

   RWTH Aachen University

   Affiliation as printed

   Faculty of Electrical Engineering and Information Technology , RWTH Aachen University , Germany

7. Dorit Merhof

   University of Regensburg

   Affiliation as printed

   Faculty of Informatics and Data Science , University of Regensburg , Germany