Region Attention Transformer for Medical Image Restoration

Zhiwen Yang,Haowei Chen,Ziniu Qian,Yang Zhou,Hui Zhang,Dan Zhao,Bingzheng Wei,Yan Xu

2024-07-12

Abstract:Transformer-based methods have demonstrated impressive results in medical image restoration, attributed to the multi-head self-attention (MSA) mechanism in the spatial dimension. However, the majority of existing Transformers conduct attention within fixed and coarsely partitioned regions (\text{e.g.} the entire image or fixed patches), resulting in interference from irrelevant regions and fragmentation of continuous image content. To overcome these challenges, we introduce a novel Region Attention Transformer (RAT) that utilizes a region-based multi-head self-attention mechanism (R-MSA). The R-MSA dynamically partitions the input image into non-overlapping semantic regions using the robust Segment Anything Model (SAM) and then performs self-attention within these regions. This region partitioning is more flexible and interpretable, ensuring that only pixels from similar semantic regions complement each other, thereby eliminating interference from irrelevant regions. Moreover, we introduce a focal region loss to guide our model to adaptively focus on recovering high-difficulty regions. Extensive experiments demonstrate the effectiveness of RAT in various medical image restoration tasks, including PET image synthesis, CT image denoising, and pathological image super-resolution. Code is available at \href{<a class="link-external link-https" href="https://github.com/Yaziwel/Region-Attention-Transformer-for-Medical-Image-Restoration.git" rel="external noopener nofollow">this https URL</a>}{<a class="link-external link-https" href="https://github.com/RAT" rel="external noopener nofollow">this https URL</a>}.

Image and Video Processing,Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The paper primarily addresses the challenges in the Medical Image Restoration (MedIR) task and proposes a new method called the Region Attention Transformer (RAT). Specifically, the paper aims to solve the following key issues:

1. **Limitations of Attention Mechanism**: Existing transformer-based methods typically use fixed and coarse-grained regions (such as the entire image or fixed-size image patches) for attention computation, which may lead to interference between irrelevant regions and fragmentation of continuous image content.

2. **Improving Interpretability and Flexibility**: To overcome the above limitations, the authors propose a Region-based Multi-Head Self-Attention (R-MSA) mechanism. This mechanism leverages the powerful "Segment Anything Model" (SAM) to dynamically segment the input image into non-overlapping semantic regions and performs attention computation based on these regions.

3. **Prioritizing Difficult Regions**: Considering the varying difficulty of restoring different image regions, the authors introduce a Focal Region Loss. This loss function can dynamically adjust the focus on different regions during training, prioritizing those that are more challenging to restore.

Through these improvements, the proposed RAT aims to achieve the following goals:

- Achieve better restoration results in various medical image restoration tasks, such as PET image synthesis, CT image denoising, and pathological image super-resolution reconstruction.
- Enhance the interpretability of the model by ensuring that attention operations are confined to similar semantic regions, thereby reducing interference between irrelevant regions.
- Further optimize image restoration performance by dynamically adjusting the importance of different regions.

Experimental results show that RAT achieves state-of-the-art performance in typical tasks such as PET image synthesis, CT image denoising, and pathological image super-resolution reconstruction.