SimLVSeg With Mamba 实验

SimLVSeg

SimLVSeg使基于视频的网络能够从稀疏注释的超声心动图视频中进行一致的LV分割。SimLVSeg包括具有时间掩蔽的自监督预训练,然后是为从稀疏注释中进行LV分割而定制的弱监督学习。我们展示了SimLVSeg如何在最大的2D+时间超声心动图数据集(EchoNet-Dynamics)上实现93.32%(95%CI 93.21-93.43%)的DICE评分,同时更高效,从而优于最先进的解决方案。SimLVSeg兼容两种类型的视频分割网络:2D超级图像和3D分割。为了展示我们方法的有效性,我们提供了广泛的消融研究,包括预训练设置和各种深度学习主干。我们进一步进行了分布外测试,以展示SimLVSeg在未见分布(CAMUS数据集)上的推广能力。该代码可在https://github.com/fadamsyah/SimLVSeg上公开获取。

SimLVSeg使基于视频的网络能够进行LV分段,从而增强性能和更高的时间一致性。SimLVSeg由两个训练阶段组成:具有时间掩蔽的自我监督预训练和LV分割的弱监督学习,专门设计用于解决稀疏注释(标记)超声心动图视频的挑战。

结构

3D U-Net架构

超声心动图视频由堆叠的2D图像组成。将时间轴视为第三维允许3D模型在超声心动图剪辑上分割LV。因此,3D U-Net被用作架构。如图4所示,我们使用具有残余单元的CNN 作为编码器,其具有5个级,其中级输出被传递到解码器。剩余单元包括两个Conv 2D层、两个实例规范层、两个PReLU激活函数和一个跳过连接。

数据集

EchoNet-Dynamics是最大的公开2D+Time心脏超声心动图,其显示了人体心脏的顶部四腔视图。提供了约10,030个心脏超声心动图视频,固定帧大小为112 x 112。

视频长度从28帧到1002帧不等,涵盖多个心跳周期,但只有两个被注释(ED和ES帧)。图1中给出了示例超声心动图序列

image-20240926151357434

实验

作者同时还做了多个消融实验,验证不同的主干网络对该框架性能的影响,实验数据如下:

Table 2 各种编码器主干的消融研究

如表2所示,实验表明性能对编码器主干是稳健的。

注意:这里torch记得用加cuda版本
1
2
#CUDA 11.6
pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu116

SegMamba

SegMamba也是一个分析3D超声心动图的网路,不同的是原论文中实验采用的体积数据,而不是2D图像数据加时间。因此,将SegMamba移植过来作为主干网络时,将时间作为第三个轴,即输入的帧数作为,SegMamba在原论文实验中表现同样十分优异。详细的集成后代码可以在这里找到https://github.com/52HZMercury/SimLVSeg-with-SegMamba

使用的环境

py3.8 torch 1.12

通道数的变化

首先需要注意的就是原论文中,SegMamba第一层是使用的4个通道,这里需要初始化为3个通道进行输入,这里解决之后,多半会遇到tensor对不上的问题,接着往下看

tensor对不上

这里tensor对不上的主要原因就是EchoNet-Dynamics数据集是每帧112112的,而SegMamba在原论文中是长宽高为128 * 128 * 128的输入。解决的方法就是把util代码里面的读取视频的代码输出时resize一下,改为128128,至于还有一个128,在SimLVSeg中输入时将帧改为128即可,即对应的

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
def load_video(filename: str) -> np.ndarray:
"""Loads a video from a file.
Args:
filename (str): filename of video
Returns:
A np.ndarray with dimensions (channels=3, frames, height, width). The
values will be uint8's ranging from 0 to 255.
Raises:
FileNotFoundError: Could not find `filename`
ValueError: An error occurred while reading the video
"""

if not os.path.exists(filename):
raise FileNotFoundError(filename)
capture = cv2.VideoCapture(filename)

frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
frame_width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))

# v = np.zeros((frame_count, frame_height, frame_width, 3), np.uint8)
v = np.zeros((frame_count, 128, 128, 3), np.uint8)

for count in range(frame_count):
ret, frame = capture.read()
if not ret:
raise ValueError("Failed to load frame #{} of {}.".format(count, filename))
#112*112
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

# 缩放为128
frame = cv2.resize(frame, (128,128))
v[count, :, :] = frame

return v
爆显存

在经历上面的操作后,在实验时,可能会遇到爆显存的错误(3090),这里的解决方法是,改小精度,将原来的32位改为了16位

实验结果

Table.2. DSC(Dice Similarity Coefficient)

Step DSC loss
2235 0.863 0.138
2980 0.922 0.079
4470 0.914 0.084
5689 0.923 0.078

Both SimLVSeg-3D and SimLVSeg-SI outperform the state-of-the-art,suggesting that the superior performance can be attributed to the SimLVSeg design rather than the selection of the underlying network architectures.

原作者在论文中说是SimLVSeg设计优秀所以性能较好而不是主干网络的选择,从移植SegMamba的实验结构来看确实如此。

改小参数量

具体的措施有:将通道数改为原来的1/2,并转为32位精度计算

改小参数量后

加入通道注意力模块

Squeeze-and-Excitation Networks

序言

在本文中,作者研究了网络设计的一个不同方面-通道之间的关系。引入了一个新的架构单元,称之为挤压和激发(SE)块,其目标是通过显式建模其卷积特征通道之间的相互依赖性来提高网络产生的表示的质量。为此,作者提出了一种允许网络执行特征重新校准的机制,通过该机制,网络可以学习使用全局信息来选择性地强调信息特征并抑制不太有用的特征

在传统的CNN架构中,通常采用卷积层和池化层来提取图像特征。然而,这种方法并没有明确地对特征通道之间的关系进行建模,导致一些通道对于特定任务的贡献相对较小,而其他通道则更重要。SE模块旨在解决这个问题。

方法

Fig. 1. A Squeeze-and-Excitation block.

SE构建块的结构如图1所示。

Squeeze操作

这个操作对应上图中的Fsq函数。上面提到了,要找到通道之间的关系,那就是要用到所有通道的信息。论文中简单起见,把每个通道都压缩成一个点,直接求平均,相当于就是这个通道的特征值。这个求平均的过程就是所谓的Squeeze。表达为公式就是:

zc=Fsq(uc)=1H×Wi=1Hj=1Wuci,jz_c = F_{sq}(u_c)= \frac{1}{ H × W} \sum_{i=1}^H\sum_{j=1}^Wu_c(i,j)

uc表示特征图U中的某个通道,zc表示经过压缩之后的向量Z。从公式中可以看出,每个尺寸为H*W的featrue map通道,经过全局平均之后,就只剩一个像素点。所以,一个的featrue map经过squeeze之后,就变成了一个的vector, 计作Z。

Excitation

特征图经过压缩后,第二步就是综合这些通道的特征值。论文提出,这个综合方法需要满足两个条件:1.要可以表达非线性关系 2.要可以提取处非互斥的关系 所以,论文提出了一个计算方法:

s=Fex(z,W)=σ(g(z,W))=σ(W2δ(W1z))s = F_{ex}(z,W) = \sigma(g(z,W)) = \sigma(W_2\delta(W_1z))

从上面可以看出,这个Excitation就是两个全联接层,第一个连接层把1* 1* C的vector Z变成一个1* 1* C/r的向量。然后后面跟一个ReLU激活 层。· 第二个连接层把1* 1* C/r的vector重新恢复为一个1* 1* C的向量。然后后面跟一个Sigmod 的激活层。

Scale

最后,根据图1,作为Squeeze和Excitation操作之后,还需要和原特征图U进行合并,也就是上面提到的“重新校准特征图”。这个合并方法也特别简单,就是把这个向量直接原特征图相乘。

xc~=Fscale(uc,sc)=scuc\widetilde{x_c} = F_{scale}(u_c,s_c) = s_cu_c

从上面的excitation可以得知,s是一个的向量,与一个1* 1* C的特征图相乘的话,再从上面的公式可以看出,就是按通道数一个一个的去把s中的每个元素拿出来与特征图的对 应通道进行相乘,就是一个标量乘以矩阵,所以可以理解就是一个scale的操作,放大缩小而已。讲到这里,感觉就是SE就是提取出了通道中某些特征,或者说某些注意力特征,通过这个s’也就是表示通道的权重的向量,去把这个权重乘以对应的特征图通道,来放大或缩小这个通道的作用。从这个角度上来看,确实就叫做通道注意力机制。

下面给出了SE与Inception和ResNet合并的方式。

Fig. 2.原始Inception模块(左)和SEInception模块(右)的模式

Fig. 3.原始残差模块(左)和SEResNet模块(右)的架构。

总结

总而言之,SE模块通过引入Squeeze和Excitation操作,通过自适应地学习每个通道的权重,增强了神经网络的表达能力和性能,使网络能够更有效地学习和利用特征通道之间的关系。

通道再次降低为初始的1/3,即24M的模型,同时加入通道注意力模块,即上面的压缩激励

降低为原来的1/3

image-20241031131654163

加入MSAA模块

CM-UNet: Hybrid CNN-Mamba UNet for Remote Sensing Image Semantic Segmentation

加入多尺度注意力聚合模块MSAA

加入MSAA后

提出CM-UNet框架:基于Mamba架构的CM-UNet框架,通过整合CNN和Mamba模块,能够在遥感图像语义分割任务中高效捕捉局部和全局信息。

设计CSMamba块:CSMamba块结合了通道和空间注意力机制,将Mamba模块扩展为能够处理图像长程依赖的组件,提升了特征选择和信息融合的精度。

多尺度注意力聚合模块(MSAA):引入MSAA模块,聚合编码器的多尺度特征,通过空间和通道的双重聚合提高特征表达能力,替代传统的跳跃连接,更好地支持解码器的多层次信息融合。

多输出监督机制:在解码器的不同层次引入多输出监督,确保各层次逐步细化分割图,从而提升最终分割精度。

整体结构

​ CM-UNet模型结构由ResNet编码器、多尺度注意力聚合模块(MSAA)和CSMamba解码器组成。编码器负责提取多层次特征,MSAA模块融合多尺度特征以增强表达,解码器则利用CSMamba块通过通道和空间注意力机制高效捕捉长程依赖关系,最终生成精细的分割图,并在各层解码器中加入多输出监督以优化分割结果。

Fig. 2. Framework of the proposed CM-UNet for remote sensing image semantic segmentation. (a) CM-UNet. (b) CSMamba Block. (c) 2D-SSM Module.

CNN编码器:采用ResNet结构作为编码器,用于提取多层次的特征信息。与传统UNet不同,CM-UNet的编码器使用的是多尺度特征提取,以便为后续模块提供更丰富的上下文信息。

多尺度注意力聚合模块(MSAA):在编码器和解码器之间,使用MSAA模块对多尺度特征进行聚合。这个模块通过空间和通道注意力机制对不同尺度的特征进行融合,增强特征表达能力,并取代了UNet中的跳跃连接。

Fig. 3. The Multi-Scale Attention Aggregation (MSAA) module.

CSMamba解码器:解码器采用CSMamba块,该模块结合通道和空间注意力机制,以及Mamba结构的线性时间复杂度特性,能够有效捕获图像的长程依赖关系。CSMamba解码器逐步上采样特征并生成输出分割图。同时引入多输出监督,在各层解码器中加入监督信号,确保分割图在不同层次上得到精细化生成。

参数量降为24.1M

CVPR 2024 的MemSAM

Deng X, Wu H, Zeng R, et al. MemSAM: Taming Segment Anything Model for Echocardiography Video Segmentation[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024: 9622-9631.

该论文实验结果

image-20241129162620958

在EchoNet-Dynamic数据集上

小segmamba实验结果(DiceLoss)

image-20241202155209382

DICE和IOU都比该文高,其中DICE高0.11个点,IOU高0.84

换成jaccard

image-20241217160043633

效果下降了,也可能是误差,决定再跑一遍

image-20241220194006242

好吧,效果依旧不太理想

换成tversky

image-20241223142447888

效果一般

在CAMUS数据集

重写dataset改为该数据集后,DICE和IOU都出现错误:表现为大于1。

猜想:每个病人只有十几帧,但是segmamba是按128帧设计的,因此输入的是128帧,多出的帧采用零填充,导致有用数据太少。

方案:使用每个病人原始的数据,进行重复,直到填满128帧。

并不是这个原因,是因为拼接的函数有点问题,改完之后可以正常显示数据,但是还不完善,只能使用一个batchsize

下面再来测试几种填充之后的实验差别

type1:

image-20241203154356256

type2:

image-20241203154228577

type3:

image-20241203193005161

type4:

image-20241203193037852

综合来看type2的padding效果最好,后续在type2上继续实验

将所有的消融实验做一个表格,包括

DICE IOU loss
type1 + Adaw + diceloss 0.912 0.845 0.095
type2 + Adaw + diceloss 0.930 0.870 0.077
type3 + Adaw + diceloss 0.923 0.860 0.084
type4 + Adaw + diceloss 0.926 0.865 0.085
type5 + Adaw + diceloss 0.912 0.845 0.095
type2 + Adaw + jaccard 0.932 0.873 0.129
type2 + Adaw + jaccard+ 2lr 0.935 0.878 0.126
type3 + Adaw + hccdice 0.931 0.871 0.080
type3 + Adaw + jaccard 0.932 0.873 0.135
type3 + Adaw + jaccard + 2lr 0.934 0.877 0.125
type3 + Adaw + jaccard + 3lr 0.932 0.873 0.128
type3 + Adaw + tversky + 2lr 0.930 0.870 0.076
type3 + SGD + diceloss 0.921 0.857 0.078
type3 + SGD + jaccard 0.919 0.850 0.149
type4 + Adaw + jaccard 0.933 0.875 0.127
type6 + Adaw + jaccard + 2lr 0.928 0.866 0.140

使用CAMUS数据集官方划分的训练集,验证集和测试集

全监督

image-20241212153454404

全监督4batch 4lr

image-20241223142653089

弱监督A

这里的弱监督是训练的时候只提取了两帧(舒张末期和收缩末期),在验证测试时也只测试了这两帧

image-20241212153606065

弱监督B

这个弱监督也是训练的时候只提取两帧,但是验证和测试时是全周期数据

image-20241212190536204

可以看到,效果非常的垃圾。。。

使用SAM模型划分的训练集,验证集和测试集

//todo

Echo-Pediatric

该数据需要首先进行预处理,使用split列中的数字划分训练集,验证集和测试集。这里使用的是<8的数据为训练集,==8的为验证集,其余的为测试集。

除此之外,该数据集中有部分脏数据需要进行处理,比如VolumeTracings.csv中部分X,Y数据为空,Frame为No Systolic,需要删除,以及部分数据标注时有误,只有144帧,但标注时标注了序号为第144帧,导致越界错误。

先使用了jaccard的dsc只有0.9063,iou为0.8288

image-20250109154150304

考虑到这个数据集和echodynamic类似。因此,接着使用diceloss进行测试,效果差不多,dsc为0.9045,iou为0.8256

image-20250109154249431

数据集的处理有点问题,重新跑一遍:

segmamba模型

segmamba dsc iou loss
jaccard+Adaw_4lr 0.9039 0.8246 0.1757
jaccard+Adaw_6lr 0.9041 0.8250 0.1752
jaccard+SGD 0.9084 0.8323 0.1677
diceloss+SGD 0.9068 0.8295 0.0932

3D_Unet模型

3D_Unet dsc iou loss
3D_Unet+jaccard 0.9090 0.8332 0.1688
3D_Unet+diceloss 0.9112 0.8369 0.0906

其他模型 都是Adaw+Dice

model name dsc iou loss
LightMUNet 0.8855 0.7946 0.1145
TMamba3D / / /
SegFormer3D 0.9123 0.8388 0.0876
UKAN3D 0.9091 0.8334 0.0910
UNet_IDC3D

LightM-UNet

运行环境版本:

LightM-UNet,这是一个基于Mamba的轻量级U形分割模型,它在显著降低参数和计算成本的同时实现了最先进的性能,参数大大减小

image-20241112173154849

image-20241114135927092

所提出的LightM-UNet的总体架构如图所示。给定一个输入图像 ,其中 、、 和 分别表示3D医疗图像的通道数、高度、宽度和切片数

image-20241114140118136

LightM-UNet首先使用深度可分卷积(DWConv)层进行浅层特征提取,生成浅层特征图。随后,LightM-UNet结合三个连续的编码器块(Encoder Blocks)从图像中提取深层特征。在每个编码器块之后,特征图中的通道数翻倍,而分辨率减半。

Encoder Block

感觉encoder都差不多

为了最小化参数数量和计算成本,LightM-UNet采用了仅包含Mamba结构的编码器块来从图像中提取深层特征。

具体来说,给定一个特征图 FlRC×H×W×DF^l \in R^{C \times H \times W \times D },其中C=32×2lC = 32 \times2^lH=H/2lH = H / 2^lW=W/2lW = W / 2^lD=D/2lD = D / 2^l,以及,编码器块首先将特征图展平并转置成的形状,其中L=H×W×DL = H \times W \times D

随后,编码器块使用个连续的RVM层来捕捉全局信息,在最后一个RVM层中通道数增加。此后,编码器块重新调整并转置特征图的形状为 (C×2,H,W,D)(C \times 2,H,W,D),紧接着进行最大池化操作以降低特征图的分辨率。

最终,第l个编码器块输出新的特征图Fl+1F^{l+1},其形状为(C×2,H/2,W/2,D/2)(C \times 2,H/2,W/2,D/2)

Residual Vision Mamba Layer (RVM Layer)

LightM-UNet提出了RVM层以增强原始的SSM块,用于图像深层语义特征提取。具体来说,LightM-UNet利用先进的残差连接和调整因子进一步增强了SSM的长距离空间建模能力,几乎不引入新的参数和计算复杂性。

实验

image-20241114184622536

LightM-UNet以其极低的参数数量脱颖而出,仅使用了1.09M个参数。与nnU-Net和U-Mamba相比,参数分别减少了99.14%和99.55%

加入通道注意力特征融合

image-20241119155343122

同时实验时加入mIOU评价指标

image-20241128123310222


SimLVSeg With Mamba 实验
http://example.com/2024/09/26/SimLVSegWithMamba/
作者
Mercury
发布于
2024年9月26日
许可协议