Search arXivSearch

arXiv · 2407.11087

Restore-RWKV: Efficient and Effective Medical Image Restoration with RWKV

Abstract

Transformers have revolutionized medical image restoration, but the quadratic complexity still poses limitations for their application to high-resolution medical images. The recent advent of the Receptance Weighted Key Value (RWKV) model in the natural language processing field has attracted much attention due to its ability to process long sequences efficiently. To leverage its advanced design, we propose Restore-RWKV, the first RWKV-based model for medical image restoration. Since the original RWKV model is designed for 1D sequences, we make two necessary modifications for modeling spatial relations in 2D medical images. First, we present a recurrent WKV (Re-WKV) attention mechanism that captures global dependencies with linear computational complexity. Re-WKV incorporates bidirectional attention as basic for a global receptive field and recurrent attention to effectively model 2D dependencies from various scan directions. Second, we develop an omnidirectional token shift (Omni-Shift) layer that enhances local dependencies by shifting tokens from all directions and across a wide context range. These adaptations make the proposed Restore-RWKV an efficient and effective model for medical image restoration. Even a lightweight variant of Restore-RWKV, with only 1.16 million parameters, achieves comparable or even superior results compared to existing state-of-the-art (SOTA) methods. Extensive experiments demonstrate that the resulting Restore-RWKV achieves SOTA performance across a range of medical image restoration tasks, including PET image synthesis, CT image denoising, MRI image super-resolution, and all-in-one medical image restoration. Code is available at: https://github.com/Yaziwel/Restore-RWKV.

Explore related subjects

Keep this discovery

Explore connections, maps & timelines

BibTeXRIS

Zhiwen Yang, Jiayin Li, Hui Zhang, Dan Zhao, Bingzheng Wei, Yan Xu. 2025-01-06. Restore-RWKV: Efficient and Effective Medical Image Restoration with RWKV. https://arxiv.org/abs/2407.11087

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

Dipole-lets: a new multiscale decomposition for MR phase and quantitative susceptibility mapping

Nondipolar phase contributions can generate severe streaking artifacts during quantitative susceptibility mapping (QSM) inversion. We propose Dipole-lets, a dipole-adapted multiscale decomposition designed to identify phase components associated with streaking artifacts before and during susceptibility inversion. Dipole- lets combine an undecimated radial decomposition with an angular partition based on proximity to the magic cone. The resulting coefficients emphasize phase components near the magic cone, where nondipolar contributions may become relatively prominent. This information was incorporated into QSM reconstruction through a data- driven fidelity weight and a Dipole-let-based regularizer that models an additional nondipolar phase component. The methods were evaluated using modified simulated QSM data and an in vivo dataset. The data-driven weighting reduced streaking artifacts while preserving anatomical detail and improved quantitative reconstruction metrics compared with magnitude-weighted TV. The Dipole-let regularizer achieved streaking suppression comparable to L1-QSM and lower reconstruction error in the modified simulated dataset. In vivo experiments further demonstrated the applicability of the proposed approach to data with strong susceptibility-induced phase perturbations. Dipole-lets provide a multiscale representation for characterizing phase components associated with streaking artifacts and incorporating this information into QSM reconstruction, reducing streaking while preserving relevant susceptibility structures, and providing a flexible basis for future QSM reconstruction methods.

eess.IV

MTMed3D: A Multi-Task Transformer-Based Model for 3D Medical Imaging

In the field of medical imaging, AI-assisted techniques such as object detection, segmentation, and classification are widely employed to alleviate the workload of physicians and doctors. However, single-task models are predominantly used, overlooking the shared information across tasks. This oversight leads to inefficiencies in real-life applications. In this work, we propose MTMed3D, a novel end-to-end Multi-task Transformer-based model to address the limitations of single-task models by jointly performing 3D detection, segmentation, and classification in medical imaging. Our model uses a Transformer as the shared encoder to generate multi-scale features, followed by CNN-based task-specific decoders. The proposed framework was evaluated on the BraTS 2018 and 2019 datasets, achieving promising results across all three tasks, especially in detection, where our method achieves better results than prior works. Additionally, we compare our multi-task model with equivalent single-task variants trained separately. Our multi-task model significantly reduces computational costs and achieves faster inference speed while maintaining comparable performance to the single-task models, highlighting its efficiency advantage. To the best of our knowledge, this is the first work to leverage Transformers for multi-task learning that simultaneously covers detection, segmentation, and classification tasks in 3D medical imaging, presenting its potential to enhance diagnostic processes. The code is available at https://github.com/fanlimua/MTMed3D.git.

eess.IV

VideoPulse: Neonatal heart rate and peripheral capillary oxygen saturation (SpO2) estimation from contact free video

Remote photoplethysmography (rPPG) enables contact free monitoring of vital signs and is especially valuable for neonates, since conventional methods often require sustained skin contact with adhesive probes that can irritate fragile skin and increase infection control burden. We present VideoPulse, a neonatal dataset and an end to end pipeline that estimates neonatal heart rate and peripheral capillary oxygen saturation (SpO2) from facial video. VideoPulse contains 157 recordings totaling 2.6 hours from 52 neonates with diverse face orientations. Our pipeline performs face alignment and artifact aware supervision using denoised pulse oximeter signals, then applies 3D CNN backbones for heart rate and SpO2 regression with label distribution smoothing and weighted regression for SpO2. Predictions are produced in 2 second windows. On the NBHR neonatal dataset, we obtain heart rate MAE 2.97 bpm using 2 second windows (2.80 bpm at 6 second windows) and SpO2 MAE 1.69 percent. Under cross dataset evaluation, the NBHR trained heart rate model attains 5.34 bpm MAE on VideoPulse, and fine tuning an NBHR pretrained SpO2 model on VideoPulse yields MAE 1.68 percent. These results indicate that short unaligned neonatal video segments can support accurate heart rate and SpO2 estimation, enabling low cost non invasive monitoring in neonatal intensive care.

eess.IV