点击下方卡片,关注 「AI视界引擎」 公众号
「AI视界引擎交流群」
三维成像中的高效牙齿分割对于正畸诊断至关重要,但由于CBCT图像中的噪声、低对比度和伪影,这仍然是一个挑战。
卷积神经网络(CNNs)和变压器(transformers)都已成为图像分割中流行的架构。然而,由于内在的局部性或计算复杂性,它们在处理长距离依赖方面的有效性受到限制。
为了解决这个问题,我们提出了T-Mamba,将共享位置编码和基于频率的特征整合到视觉Mamba中,以解决空间位置保留和频率域特征增强的限制。
此外,我们还设计了一个门选择单元,以自适应地整合空间域中的两个特征和频率域中的一个特征。T-Mamba是首次将基于频率的特征引入视觉Mamba的工作。
大量实验证明,T-Mamba在公共牙齿CBCT数据集上取得了新的SOTA结果,并且相较于之前的方法在较大幅度上取得了性能提升,即 IoU + 3.63%,SO + 2.43%,DSC + 2.30%,HD - 4.39mm,ASSD - 0.37mm。
现代数字牙科的进化关键在于三维(3D)成像的获取与分割。这项技术在口腔及颌面学科领域有多种用途,包括正畸诊断和治疗规划。牙齿分割,即从3D数字模型中提取具有强度和密度的一组体素,是数字工作流程中的必要阶段。然而,精确的牙齿分割过程由于多种原因而具有挑战性。首先,由于3D数字模型固有的高噪声和低对比度特性,难以观察。此外,CBCT图像中金属充填物和修复体的存在引入了导致畸变的伪影,显著增加了识别牙齿的难度。最后,CBCT通常在自然咬合状态下获取,由于上下牙齿密度相似,区分它们也构成了挑战。
深度学习在医学图像分割领域得到了广泛应用,许多研究者致力于利用深度学习技术实现CBCT图像中牙齿的自动分割,如Chen等人(2023年);Cui等人(2019年);Cui等人(2021年);Yin等人(2023年)。卷积神经网络(CNNs)和Transformers架构因其能够学习复杂的图像特征并提供准确的分割结果,在医学图像分割方面显示出巨大的潜力。CNNs能够捕获平移不变性并提取局部特征,而Transformers擅长捕捉全局上下文信息并提高长距离依赖的能力。由于这种互补特性,许多研究探索了通过混合网络架构将Transformers融入CNNs中。然而,Transformers的明显缺点是资源密集型,因为自注意力机制与输入大小成二次方增长,并且在处理高分辨率生物医学图像时在速度和内存使用方面构成挑战。尽管大量努力致力于降低Transformers的计算复杂性,但这往往以牺牲模型准确性为代价,如Maaz等人(2022年);Zhai等人(2023年)。因此,在CNNs中有效增强长距离依赖仍然是一个未解决的问题。
近年来,受到Mamba在语言建模领域的成功启发,许多研究将这种成功从语言领域转移到视觉领域,旨在实现线性复杂度同时不牺牲全局感受野 Liu等人(2024);Zhu等人(2024)。然而,我们发现在利用视觉Mamba帮助卷积神经网络(CNN)建模长距离依赖时存在两个局限性。首先,CNN可以直接处理2D或3D特征,而视觉Mamba的结构旨在处理1D特征序列。在CNN和视觉Mamba混合架构中,视觉Mamba的特征转换不可避免地导致空间位置信息的丢失。这种空间上下文在需要密集精确位置预测的任务中尤为重要。其次,由于CT和X射线等医学图像的固有成像原理,这些图像从视觉角度具有高噪声和低对比度的自然属性。对于这类图像,基于频域的特征表示更准确、独特和鲁棒。然而,CNN和视觉Mamba模型通常仅从空域提取语义特征,忽略了丰富的基于频域的信息。正如Azad等人(2021)所指出的,高频成分捕捉纹理细节,而低频成分编码形状信息。因此,将频域特征与空域表示相结合,有望增强医学图像中的图像特征提取,从而提高预测的准确性。
受到这两个局限性的启发,我们设计了一个名为T-Mamba的网络,它将我们提出的Tim(牙科视觉Mamba)块与DenseVNet Gibson等人(2018)的多尺度特征相结合,用于牙科CBCT图像分割。Tim块展现出三个关键优势:(1)它包含一个共享的双位置嵌入,用于补偿在重塑操作期间丢失的位置信息。值得注意的是,我们在每个特征尺度内使用单一位置嵌入。这不仅在相同尺度的特征图之间保证了空间位置的保留,还减少了模型参数和计算负担。(2)它提取频域中的图像特征,因此我们可以为具有高噪声和低对比度的医学图像导出更准确、独特和鲁棒的特征表示。基于不同尺度特征的性质,我们为每个尺度定制了不同的带通滤波策略。(3)它包括一个门选择单元,用于自适应地整合空域中的两个特征(正向和反向)和频域中的一个特征。
门选择单元是数据依赖的,可以根据输入序列标记为三个不同的特征分配权重。
在牙齿CBCT分割任务上进行了大量实验,以验证T-Mamba的有效性。我们的T-Mamba在公共牙齿CBCT数据集上的表现大幅超越了先前最先进的结果,即IoU提高了3.63%,SO提高了2.43%,DSC提高了2.30%,HD降低了4.39毫米,ASSD降低了0.37毫米。此外,我们还进行了消融研究,以验证我们提出的Tim块中三个组件的有效性。据我们所知,T-Mamba是首次将频域特征融入到视觉Mamba框架中的开创性工作。
主要贡献可以总结如下:
- 我们提出了T-Mamba,它将我们设计的Tim块与DenseVNet结合,用于牙齿CBCT分割的全局和局部视觉上下文建模。
- Tim块是首次通过引入基于频率的特征,为高噪声和低对比度的医学图像提取更鲁棒和独特的表示。
- 在没有额外修饰的情况下,T-Mamba在公共牙齿CBCT数据集上取得了新的SOTA结果,并且以较大的幅度超过了先前SOTA方法,如IoU提高了3.63%,SO提高了2.43%,DSC提高了2.30%,HD降低了4.39毫米,ASSD降低了0.37毫米。
3 方法论
3.1 实验设计
实验设计基于随机完全区组设计(RCBD),设有三个重复。
Preliminaries for Mamba
高级状态空间模型(SSM),即结构化状态空间序列模型(S4)和Mamba,是一类系统,它们通过隐藏状态 将一维连续函数或序列 进行映射。从数学上讲,这些模型通常被表述为线性常微分方程(ODEs),如公式(1)所示:
\begin{split} h'(t) &= \mathrm{A}h(t) + \mathrm{B}x(t), \\ y(t) &= \mathrm{C}h(t), \end{split} \tag{1}
其中参数包括 作为演化参数,以及 、 作为投影参数。
基于SSM的模型作为连续时间模型,在整合到深度学习算法中时应当进行离散化。这种离散化转换对于将模型与输入数据中蕴含的底层信号的采样率对齐至关重要[1]。给定输入 ,一个在信号流程中的长度为 的采样向量[2],公式(1)可以使用零阶保持(ZOH)规则离散化为如下形式:
\begin{split} h\_t &= \overline{\mathbf{A}}h\_{t-1} + \overline{\mathbf{B}}x\_t \\ y\_t &= \overline{\mathbf{C}}h\_t \\ \overline{\mathbf{A}} &= e^{\Delta\mathbf{A}} \\ \overline{\mathbf{B}} &= \Delta\mathbf{A}^{-1}(e^{\Delta\mathbf{A}} - \mathbf{I}) \cdot \Delta\mathbf{B} \\ \overline{\mathbf{C}} &= \mathbf{C}, \end{split} \tag{2}
其中 是时间尺度参数。
最终,模型通过在结构化卷积核 内的全局卷积操作计算输出 :
\begin{split} \overline{\mathbf{K}} &= (\mathbf{C}\overline{\mathbf{B}}, \mathbf{C}\overline{\mathbf{A}}\overline{\mathbf{B}}, \mathbf{C}\overline{\mathbf{A}}^2\overline{\mathbf{B}}, ..., \mathbf{C}\overline{\mathbf{A}}^{L-1}\overline{\mathbf{B}}) \\ y &= \overline{\mathbf{K}} \otimes y\_t. \end{split} \tag{3}
我们通过利用Mamba的线性扩展优势,增强了卷积神经网络(CNN)的表征建模能力,并提出了一种用于图像分割的通用网络T-Mamba。T-Mamba的网络架构如图1所示,它将我们提出的Tim模块与DenseVNet在多尺度特征上进行了融合。T-Mamba包含三个特征尺度,并采用了经典的单阶段V型架构。我们在DenseVNet的每个CNN层后简单地插入Tim模块,以建模长距离依赖。T-Mamba能够在空间域和频率域捕获局部的细粒度特征和长距离依赖。我们希望它能为生物医学图像分析中高效的长距离依赖建模开辟新的途径。
原始的Mamba块是为1-D序列设计的,这不适用于需要空间感知理解的任务。Vision Mamba Zhu等人,2024年提出的Vim块在Mamba块上集成了双向序列建模。基于这个强大的设计,我们通过在Vim块中引入三个组件来进一步增强视觉特征建模,这对于2D和3D牙齿分割任务至关重要:1) 共享双向位置编码补偿;2) 基于频率的带通滤波;3) 门控选择单元。我们的牙齿视觉Mamba(Tim)块如图1所示。
共享双向位置编码补偿。 我们的T-Mamba网络将卷积层的局部特征提取能力与SSMs捕捉长距离依赖的能力相结合。卷积层通常处理2-D或3-D特征图,而我们设计的Tim块专注于处理1-D序列。因此,将高维特征重塑为1-D特征标记是必要的。然而,这一过程不可避免地会导致关键位置信息的丢失,这对于密集预测任务是至关重要的。为了减轻这一点,我们采用共享位置嵌入来补偿在重塑操作中丢失的位置信息。
具体来说,对于一个形状为(B, C, H, W, D)的输入特征,我们首先将其展平为1-D特征标记,其形状为(B, L, C),其中L = H × W × D,然后向特征标记中添加一个可学习的位置嵌入,其形状为(C, L),以保留位置信息。遵循[transformer is all u need],我们通过正弦函数初始化位置嵌入:
\begin{split} PE\_{(pos,2i)}=sin(pos/10000^{2i/d\_{\text{nodel}}}), \\ PE\_{(pos,2i+1)}=cos(pos/10000^{2i/d\_{\text{nodel}}}),\end{split} \tag{4}
其中pos是沿L的位置,i是沿C的索引。关于Tim块的输出,我们还需要将1-D特征标记重塑为高维特征,以供下一次卷积操作使用。为了进一步强化1-D特征标记中的空间信息,在将其重塑回高维特征之前,之前使用过的位置嵌入被再次添加到1-D特征标记中。请注意,在每个特征尺度上我们只使用单一的位置嵌入。这种做法确保了在同一尺度特征图中空间位置保持不变,同时减少了模型参数和计算负担。我们认为,在Tim块的输入和输出中添加共享双向位置嵌入,可以显著保留高维特征的位置信息。通过消融实验,我们验证了使用共享双向位置编码相比于使用单一位置嵌入能带来更高的性能。
基于频率的带通滤波。傅里叶域在提取基于频率的图像信息分析中起着主要作用,很明显,即使是在质量较差的X射线和CT图像中,也可以在频率域中捕捉到细微的细节和轮廓。Said等人,2004;余毅等人,2021年和Azad等人,2021;李等人,2024年。卷积层具有强烈的纹理诱导偏差,它们倾向于学习基于纹理的特征。以频率的形式表示对象可以减少纹理偏差的影响,因为只有高频部分负责纹理信息(如边界),而低频部分可能与形状有关。受到这一点的启发,我们通过增强频率域中的特征提取来演进Vim的结构。具体来说,我们首先将序列标记X转换到傅里叶域,然后使用可学习的权重参数提取频率特征,并实施带通滤波,最后进行逆傅里叶变换以恢复信号。之后,通过带有最大池化操作的激活来聚合频率特征。整个过程可以表述为:
$\begin{split}\mathcal{F}_{freq}=IFFT(Bandpass(W_{f}(FFT(X))))* Maxpool(Z^{\prime}),\ Bandpass=\begin{cases}&X*(|X|<s_{low}),x\in low-level="" features,\="" &x*(s_{low}<|x|
其中、是带通滤波的阈值。在实验中,我们设置=0.1和=0.9。低级别、中级别、高级别分别表示我们网络中的三种不同特征尺度。
门选择单元。演进的Vim块包括空间域中的两个特征(前后方向)和频率域中的一个特征。我们设计了一个门选择单元,旨在自适应地融合这些独特的特征。输入嵌入序列首先下采样到一个固定维度,如2048,并通过MLP(多线性投影)进行投影,并通过一个全连接层预测与三个特征相对应的三个比例。之后,通过三个特征的加权和得到,然后通过一个线性层进行投影,门选择单元输出的总和、残差信息以及共享位置嵌入。
\begin{split}& S\_{forward},S\_{backward},S\_{freq}=FC(MLP(Adaptive\ Pooling(X))),\\ & f\_{fuse}=(S\_{forward}*F\_{forward}+S\_{backward}*F\_{backward}+S\_ {freq}*F\_{freq}),\\ & O\_{gate}=f\_{fuse}+X+shared\ pos,\end{split} \tag{6}
请注意,内嵌公式和单独成行的公式未翻译,并且内嵌公式与文本之间保留了空格。
其中,shared pos 是前面提到的共享位置嵌入。门选择单元(Gate Selection Unit)是数据依赖的,因为这三个权重系数是从源数据 X 计算得出的,然后这些权重系数被用来更新 X 的三种不同形式的特征。因此,门选择单元可以自适应地根据不同的输入调整这三种形式特征的组合,从而获得更好的特征表示。
4 Experiments and Discussion
我们进行了与一系列现有最佳(SOTA)方法和众多视觉Mamba变体的广泛比较,展示了我们T-Mamba的优越性,并取得了新的最佳成果。此外,我们还对Tim块中的每个组件进行了消融研究。
4.1 实验设置
实验协议的设置如下:样本量在所有实验中固定为 n = 20。每次试验的持续时间设置为120秒。试验间间隔保持恒定的60秒。采用被试内设计,每位参与者经历所有实验条件。因变量为反应时间,以毫秒为单位测量。自变量是刺激类型(视觉或听觉)和难度水平(简单、中等或困难)。
4.1.1 3D CBCT Dataset
在我们的研究中使用的3D CBCT数据集是由Cui等人于2022年发布的大规模CBCT数据集中收集而来的。这个大规模数据集用于个体牙齿和牙槽骨的分割和重建,它由中国15个不同中心的4938个CBCT扫描组成,数据分布各不相同。然而,由于医院隐私问题和法规政策,只发布了一部分数据。我们研究中的数据设置与钟等人于2024年的设置相同,总共使用了129个扫描,分为训练集103个扫描和测试集26个扫描。这些扫描的物理分辨率是各向同性的,从0.2到0.4 不等。部分样本展示在图2中。
4.1.2 Implementation Details
所有实验都是在配备有32G内存的NVIDIA V100 GPU上使用pyTorch进行的。图像预处理和数据增强严格遵循钟等人的方法(2024年)。对于3D CBCT牙齿数据集,我们将每个3D图像重新采样到统一的体素间隔0.5 × 0.5 × 0.5mm,然后随机裁剪每张图像至160 × 160 × 96的大小作为模型输入。训练过程采用了AdamW优化器,betas设置为(0.9,0.999),动量为0.8,使用ReduceLROnPlateau学习率调度器,初始学习率为5e-3,权重衰减为5e-5。遵循钟等人(2024年)的配置,T-Mamba从零开始训练20个周期。采用了五种评估指标来评估网络的性能,包括Dice相似系数(DSC)、交并比(IoU)、平均交并比(mIoU)、准确率(ACC)、Hausdorff距离(HD)、平均对称表面距离(ASSD)以及表面重叠(SO)。
Quantitative Comparisons
为了评估我们的T-Mamba在3D CBCT牙齿数据集上的性能,我们与多种神经网络架构的几种最先进(SOTA)3D网络进行了全面比较。这些网络包括UNet3D及其变体(MultiRe-sUNet3D、AttentionUNet3D和PMFSNet3D),DenseVNet及其变体(DenseVoxelNet),基于变换器的网络如UNETR、SwinUNETR、TransBTS、nnFormer和3D UX-Net。
表2展示了在计算复杂性和有效性方面的比较。总体而言,T-Mamba在所有指标上均优于当前的最新方法,并且从准确度角度来看,大幅优于SOTA方法PMFSNet3D。具体来说,T-Mamba将Hausdorff距离(HD)降低了4.39毫米,平均对称表面距离(ASSD)降低了0.37毫米。此外,T-Mamba将交并比(IoU)提升了3.63%,相似度重叠(SO)提升了2.43%,以及Dice相似系数(DSC)提升了2.30%。
我们通过逐个添加的方式进行了实验,以展示我们在Tim块中提出的各个组件的有效性。首先,我们直接将普通的Vim块添加到DenseVNet中,以引入具有线性复杂度的全局特征关系建模。度量指标IoU、SO、DSC分别提升了2.79%、0.93%、1.89%;但HD和ASSD指标分别增加了1.32mm和0.11mm。这一现象表明,能够从图像特征中捕捉长距离依赖关系的Vim块对牙科CBCT分割有益,但仍然存在一些不足。
为了进一步提升Vim块的性能,我们引入了三个组件:
- 共享双向位置编码补偿;
- 基于频率的带通滤波;
- 门控选择单元。
关于位置编码补偿,当仅使用在Tim块中的归一化操作之前插入的前位置编码时,除HD外的所有指标都会变差。另一方面,当在门控选择单元中使用后位置编码时,只有ASSD和SO指标有所提升。令人惊讶的是,共享双向位置编码能够在所有指标上带来显著的改进,尤其是对于HD和ASSD指标。这三个实验表明,共享双向位置编码对于补偿输入和输出重塑操作中位置信息的丢失至关重要,可以在不引入额外模型参数的情况下提高分割准确性。之后,我们添加了从输入序列到输出序列的残差连接,这一步骤可以在所有指标上带来轻微的提升。
我们还利用基于频率的带通滤波在频率域中提取独特的特征表示。IoU和DSC指标有小幅上升。最终,添加门控选择单元以融合三种类型的特征,这个单元显著提升了IoU、SO、DSC指标。总之,我们提出的Tim块与DenseVNet基线相比,分别提升了IoU、SO、DSC指标3.74%、2.65%、2.45%,并将HD和ASSD指标分别降低了7.03mm和0.72mm。此外,我们的Tim块与Vim块相比,取得了更高的结果。
结论
我们提出了T-Mamba,以实现对牙科CBCT分割的全局和局部视觉上下文建模。得益于提出的Tim块,它将共享位置编码和基于频率的特征集成到视觉Mamba中,我们解决了医学图像在空间位置保持和频域特征增强方面的局限性,这些图像具有高噪声和低对比度。广泛的实验表明,T-Mamba在公共牙科CBCT数据集上取得了新的SOTA结果,显示Tim在生物医学图像分析中高效地进行长距离依赖建模具有巨大潜力。
点击上方卡片,关注 「AI视界引擎」 公众号