HyperAIHyperAI

Command Palette

Search for a command to run...

4 months ago

RWKV-UNet: Improving UNet with Long-Range Cooperation for Effective Medical Image Segmentation

Jiang Juntao ; Zhang Jiangning ; Liu Weixuan ; Gao Muxuan ; Hu Xiaobin ; Yan Xiaoxiao ; Huang Feiyue ; Liu Yong

RWKV-UNet: Improving UNet with Long-Range Cooperation for Effective
  Medical Image Segmentation

Abstract

In recent years, there have been significant advancements in deep learningfor medical image analysis, especially with convolutional neural networks(CNNs) and transformer models. However, CNNs face limitations in capturinglong-range dependencies while transformers suffer high computationalcomplexities. To address this, we propose RWKV-UNet, a novel model thatintegrates the RWKV (Receptance Weighted Key Value) structure into the U-Netarchitecture. This integration enhances the model's ability to capturelong-range dependencies and improve contextual understanding, which is crucialfor accurate medical image segmentation. We build a strong encoder withdeveloped inverted residual RWKV (IR-RWKV) blocks combining CNNs and RWKVs. Wealso propose a Cross-Channel Mix (CCM) module to improve skip connections withmulti-scale feature fusion, achieving global channel information integration.Experiments on benchmark datasets, including Synapse, ACDC, BUSI, CVC-ClinicDB,CVC-ColonDB, Kvasir-SEG, ISIC 2017 and GLAS show that RWKV-UNet achievesstate-of-the-art performance on various types of medical image segmentation.Additionally, smaller variants, RWKV-UNet-S and RWKV-UNet-T, balance accuracyand computational efficiency, making them suitable for broader clinicalapplications.

Code Repositories

juntaojianggavin/rwkv-unet
Official
pytorch
Mentioned in GitHub

Benchmarks

BenchmarkMethodologyMetrics
medical-image-segmentation-on-acdcRWKV-UNet
Dice Score: 0.9217
medical-image-segmentation-on-miccai-2015-1RWKV-UNet
Avg DSC: 84.02
Avg HD: 15.7

Build AI with AI

From idea to launch — accelerate your AI development with free AI co-coding, out-of-the-box environment and best price of GPUs.

AI Co-coding
Ready-to-use GPUs
Best Pricing
Get Started

Hyper Newsletters

Subscribe to our latest updates
We will deliver the latest updates of the week to your inbox at nine o'clock every Monday morning
Powered by MailChimp