医图论文 ACM MM‘24 | BSBP-RWKV：用于高效医学图像分割的背景抑制与边界保留方法-CSDN博客

点击下方“ReadingPapers”卡片，每天获取顶刊论文解读

论文信息

题目：BSBP-RWKV: Background Suppression with Boundary Preservation for Efficient Medical Image Segmentation

BSBP-RWKV：用于高效医学图像分割的背景抑制与边界保留方法

作者：Xudong Zhou、Tianxiang Chen

论文创新点

首次应用RWKV于医学图像任务：作者首次成功将RWKV应用于医学图像任务，为后续基于RWKV的高效准确方法的发展提供了新的基准和有价值的见解，开拓了医学图像分割研究的新方向。
设计DWT - PMD RWKV模块：结合Perona - Malik扩散（PMD）在抑制噪声且保留边界细节的优势与RWKV的高效结构，设计了DWT - PMD RWKV模块。该模块通过离散小波变换（DWT）与PMD的结合，在抑制背景噪声干扰的同时，能够有效保留病变区域的边界特征，提升分割的准确性。
提出多步龙格 - 库塔卷积模块：鉴于仅靠RWKV设计难以充分提取局部特征，作者提出多步龙格 - 库塔卷积模块。
设计形状优化损失函数：针对现有医学图像分割损失函数在空间域优化的局限性，提出结合频率信息的形状优化损失函数。

摘要

医学图像分割对于疾病诊断和治疗规划具有重要意义。尽管取得了多项进展，但目前大多数方法存在两个问题：一是对抑制影响分割精度的背景噪声干扰重视不足；二是效率不够高，尤其是在处理高分辨率图像时。为应对这两个挑战，作者借助一种传统去噪方法和一种新型高效网络结构，提出了BSBP - RWKV模型，用于实现精确且高效的医学图像分割。具体而言，作者结合了Perona - Malik扩散（PMD）在抑制噪声的同时保留边界细节的优势，以及RWKV高效的结构特点，在编码器的一个分支中设计了DWT - PMD RWKV模块。该模块能够在抑制背景噪声干扰的同时，保留病变区域的边界细节。然后，作者将去噪后的病变边界线索输入到所提出的多步龙格 - 库塔卷积模块中，以补充更多局部细节。此外，作者还提出了一种新颖的形状优化损失函数，该函数可以在空间域和频率域中，使预测的病变区域形状与真实标签掩码对齐。在ISIC 2016和Kvasir - SEG数据集上的实验表明，BSBP - RWKV模型在精度和效率方面表现卓越。具体来说，与当前最优方法相比，BSBP - RWKV模型的复杂度降低了5.8倍，并且在推理过程中，对于每一幅1024×1024的图像，GPU内存使用量减少了62.7%以上。

3. 方法

3.1 概述

作者提出了BSBP - RWKV，这是一种基于PMD和多步龙格 - 库塔方法的医学图像分割框架。如图3所示，编码器采用双分支结构，由级联的DWT - PMD RWKV模块特征提取分支和并行的多步龙格 - 库塔模块分支组成。DWT - PMD RWKV模块特征提取分支用于抑制背景噪声干扰，同时保留病变目标的边界，有助于形状感知分割。多步龙格 - 库塔模块分支旨在将DWT - PMD RWKV模块分支的边界预测与具有准确位置线索的目标主体特征相结合，进一步优化形状。共享主干用于为第一个DWT - PMD RWKV模块获取初始输入，并为第一个多步龙格 - 库塔模块提供部分输入。每个阶段的第一个多步龙格 - 库塔模块的输入由DWT - PMD RWKV模块的输出和上一阶段多步龙格 - 库塔模块的输出组成。生成的特征图随后被输入到普通解码器中，通过四个阶段逐步放大融合后的特征，直到到达分割头生成最终的掩码结果。作者设计了一个特定的损失函数来监督最终的预测。

3.2 DWT - PMD RWKV模块

Perona - Malik扩散（PMD）最初用于图像去噪任务。它可以通过保留图像边界和抑制噪声干扰来提高图像质量。医学图像经常受到背景噪声干扰的影响，有时病变区域边界模糊，这给准确的形状感知医学分割带来了巨大挑战。因此，作者打算构建一个受PMD启发的RWKV模块，作用于特征图，以便在过滤背景干扰的同时，仍能保留一些目标边界线索。

给定一个输入特征图，其PMD方程为：

其中是扩散系数；是扩散步长，可以看作是特征图所在的层深度；是一个控制扩散程度的正常数，在作者的实验中默认设置为1。值得注意的是，方程（1）是一个各向异性扩散方程：在梯度幅度较小的平坦或平滑区域（），扩散系数较大，这意味着扩散作用强，方程（1）起到高斯平滑的作用，去除噪声干扰；在靠近目标边界的地方，梯度幅度较大（），系数接近零，这意味着扩散作用弱，因此可以保留边界细节。

方程（1）也可以改写为以下形式：

其中和分别表示特征图在水平和垂直方向上的梯度。

另一方面，输入特征图的离散小波变换（DWT）可以表示为：

其中表示低频分量，主要反映图像中目标的基本结构。、和分别表示图像在水平、垂直和对角方向上的高频分量，主要捕捉边界细节。通过用近似，用近似，并将扩散步长设置为1，可以将方程（2）转换为离散形式：

在通过PMD增强特征图后，作者将扩散输出输入到由[13]实现的RWKV层中。通过在一个编码器分支的所有层中堆叠多个DWT - PMD RWKV模块（如图4所示），作者的BSBP - RWKV具备了在保留病变区域边界特征的同时抑制背景噪声干扰的能力。

3.3 多步龙格 - 库塔模块

基于ODE的方法已被证明在分割任务中是有效的。然而，大多数方法基于单步ODE求解器，由于仅使用前一步的信息来进行下一步的预测，不可避免地会导致一定程度的目标特征损失。受多步龙格 - 库塔方法的启发，作者提出了多步龙格 - 库塔模块（如图5所示），该模块输入并整合DWT - PMD RWKV模块的边界输出和目标主体位置特征，以进一步优化形状感知分割的质量。

多步龙格 - 库塔方法优于ResNet所基于的欧拉方法和经典龙格 - 库塔方法，它是一种多步ODE求解器，仅通过前两步的预测就能实现三阶预测精度。它不仅可以利用先前的值进行更精细的近似，而且有研究表明它优于其他常用的多步方法，如三阶Adams - Bashforth方法。多步龙格 - 库塔方法的公式如下：

其中，

3.4 形状优化损失

医学图像分割是一个像素级的分类问题，旨在将图像中的每个像素准确分类为目标或背景。通常，这个问题通过交叉熵损失和Dice损失来解决，其中Dice损失由Carole H Sudre等人提出，用于医学图像分割任务，以解决类别不平衡问题。然而，现有的医学图像分割损失函数是在空间域中定义的。当预测的病变区域经过迭代优化，在形状上与真实标签相似时，在空间域中进一步优化模型就变得具有挑战性，这可能导致模型陷入局部最优解。

为了解决上述问题，作者提出了形状优化损失，它通过将基于离散小波变换（DWT）的特定频率损失与现有的空间损失相结合。当预测掩码和真实标签在空间域中变得相似，模型难以进一步优化时，特别是在病变区域的边界处，该损失可以利用频率线索突出它们之间的差异，帮助模型跳出局部最优解。这是因为频率域线索对目标和背景边界处的梯度差异更为敏感。作者提出的定义如下：