P-Mamba: Marrying Perona Malik Diffusion with Mamba for Efficient Pediatric Echocardiographic Left Ventricular Segmentation 原文链接:https://arxiv.org/abs/2402.08506

1. 引言

超声波心动图是早期检测和治疗先天性心脏病的手段,准确分割其中的心脏结构(尤其是左心室)是关键。但其中的背景噪声会导致分割不精确,且需要考虑分割效率。

本文提出P-Mamba,可在减小噪声的同时保留局部目标边界细节,以达到最优的性能。此外,本文方法还有很高的效率。

P-Mamba的灵感来源为选择性状态空间模型Mamba和Perona–Malik扩散(PMD),前者可更好地建模长距离依赖性,且是硬件感知的、线性复杂度算法,后者则可减小图像噪声并保留边界细节。

2. 方法

如上图所示,P-Mamba分为3个组件:视觉Mamba编码器、基于DWT的PMD编码器、解码器。视觉Mamba编码器用于捕捉全局依赖并保证高效率,基于DWT的PMD编码器抑制背景噪声并保留边界细节。解码器分为分割头和全卷积头。

2.1 基于DWT的PMD块

给定输入特征图

u

u

u,其PMD公式为:

u

t

=

d

i

v

(

g

(

u

)

u

)

\frac{\partial u}{\partial t}=div(g(|\nabla u|)\nabla u)

∂t∂u​=div(g(∣∇u∣)∇u)

其中

g

(

u

)

=

1

1

+

(

u

k

)

2

g(|\nabla u|)=\frac1{1+(\frac{|\nabla u|}k)^2}

g(∣∇u∣)=1+(k∣∇u∣​)21​为扩散系数,

t

t

t为扩散步数(可视为特征图的层数),

k

k

k为控制扩散程度的正常数。该公式为各向异性扩散公式,在平坦或光滑区域,梯度值很小(

u

0

|\nabla u|\rightarrow0

∣∇u∣→0),扩散系数

g

g

g很大,扩散较强,可视为高斯平滑。对于目标边界,梯度值很大,则

g

g

g接近0,扩散较弱,可保留边界细节。可将上式重写为:

u

t

=

x

{

g

(

(

u

k

x

)

2

+

(

u

k

y

)

2

)

u

k

x

}

+

y

{

g

(

(

u

k

x

)

2

+

(

u

k

y

)

2

)

u

k

y

}

\frac{\partial u}{\partial t}=\frac{\partial }{\partial x}\left\{g\left(\sqrt{(\frac{\partial u_k}{\partial x})^2+(\frac{\partial u_k}{\partial y})^2}\right)\frac{\partial u_k}{\partial x}\right\}+\frac{\partial }{\partial y}\left\{g\left(\sqrt{(\frac{\partial u_k}{\partial x})^2+(\frac{\partial u_k}{\partial y})^2}\right)\frac{\partial u_k}{\partial y}\right\}

∂t∂u​=∂x∂​{g((∂x∂uk​​)2+(∂y∂uk​​)2

​)∂x∂uk​​}+∂y∂​{g((∂x∂uk​​)2+(∂y∂uk​​)2

​)∂y∂uk​​}

其中

u

x

\frac{\partial u}{\partial x}

∂x∂u​和

u

y

\frac{\partial u}{\partial y}

∂y∂u​为特征图的水平和垂直梯度。特征图的离散小波变换(DWT)可表达为

u

i

=

D

W

T

(

u

)

,

i

{

u

L

L

,

u

L

H

,

u

H

L

,

u

H

H

}

u_i=DWT(u),i\in\{u_{LL},u_{LH},u_{HL},u_{HH}\}

ui​=DWT(u),i∈{uLL​,uLH​,uHL​,uHH​}

其中

u

L

L

u_{LL}

uLL​为特征图的低频部分,

u

L

H

,

u

H

L

,

u

H

H

u_{LH},u_{HL},u_{HH}

uLH​,uHL​,uHH​分别为水平、垂直和对角方向的高频部分,主要包括边缘细节。

分别使用

u

L

H

u_{LH}

uLH​和

u

H

L

u_{HL}

uHL​近似

u

x

\frac{\partial u}{\partial x}

∂x∂u​和

u

y

\frac{\partial u}{\partial y}

∂y∂u​,并把扩散步长

δ

t

\delta t

δt视为1,则离散形式的PMD公式为:

u

k

=

u

k

1

+

g

[

(

u

L

H

2

+

u

H

L

2

)

u

L

H

]

L

H

+

g

[

(

u

L

H

2

+

u

H

L

2

)

u

H

L

]

H

L

u_k=u_{k-1}+g\left[(\sqrt{u_{LH}^2+u_{HL}^2})\cdot u_{LH}\right]_{LH}+g\left[(\sqrt{u_{LH}^2+u_{HL}^2})\cdot u_{HL}\right]_{HL}

uk​=uk−1​+g[(uLH2​+uHL2​

​)⋅uLH​]LH​+g[(uLH2​+uHL2​

​)⋅uHL​]HL​

使用PMD增强特征图后,将输出送入ResNet块中。堆叠多个基于DWT的PMD块,组成基于DWT的PMD编码器。

2.2 视觉Mamba块

初始阶段,2D输入

R

H

×

W

×

C

\mathbb R^{H\times W\times C}

RH×W×C被变换为拉长的patch

P

N

P_N

PN​,尺寸为

M

×

(

N

2

C

)

M\times (N^2\cdot C)

M×(N2⋅C),其中

N

N

N为patch的大小,

M

M

M为patch的数量。然后,

P

N

P_N

PN​通过线性投影变为

D

D

D维向量,并添加位置编码

E

pos

R

M

×

D

E_\text{pos}\in\mathbb R^{M\times D}

Epos​∈RM×D,得到

X

0

=

[

x

1

W

;

x

2

W

;

;

x

M

W

]

+

E

pos

X_0=[x^1W;x^2W;\cdots;x^MW]+E_\text{pos}

X0​=[x1W;x2W;⋯;xMW]+Epos​

其中

x

m

x^m

xm为

P

N

P_N

PN​的第

m

m

m个patch,

W

R

(

N

2

C

)

×

D

W\in\mathbb R^{(N^2\cdot C)\times D}

W∈R(N2⋅C)×D为可学习投影矩阵。第

l

1

l-1

l−1层的输出

X

l

1

X_{l-1}

Xl−1​会通过第

l

l

l层得到

X

l

X_l

Xl​:

X

l

=

V

i

m

(

X

l

1

)

+

X

l

1

X_l=Vim(X_{l-1})+X_{l-1}

Xl​=Vim(Xl−1​)+Xl−1​

2.3 损失函数

本文使用交叉熵损失,包含主要损失

L

p

r

i

m

L_{prim}

Lprim​(融合+分割头)和3个辅助损失

L

f

c

n

L_{fcn}

Lfcn​(融合+FCN头),

L

p

m

d

L_{pmd}

Lpmd​(基于DWT的PMD+分割头),

L

v

i

m

L_{vim}

Lvim​(视觉Mamba+分割头)。

3. 实验结果与分析

3.3 结果比较

与基于CNN和ViT的方法相比,P-Mamba能达到最优性能。此外,PMD的设计使得本文方法受到噪声的影响最小。

3.4 消融研究

将基于DWT的PMD块替换为只有边缘保留作用的Sobel算子,性能有所下降。这是因为前者还有噪声抑制作用。

将视觉Mamba替换为ViT,性能同样下降。

3.5 模型效率比较

实验表明,本文的P-Mamba比其余方法有更低的计算量和参数。PMD块并未带来过多的参数。

相关文章

评论可见,请评论后查看内容,谢谢!!!
 您阅读本篇文章共花了: