Spectral U-Net:利用双树复小波变换提高图像分割的精度和细节重建 !

图像处理机器学习关系型数据库

点击下方卡片,关注 「AI视界引擎」 公众号

( 添加时备注:方向+学校/公司+昵称/姓名 )

picture.image

picture.image

本文介绍了一种基于光谱分解的深度学习网络Spectral U-Net,通过利用双树复小波变换(DTCWT)进行下采样和反双树复小波变换(iDTCWT)进行上采样。

作者设计了对应的波段和i波段,并集成到U-Net架构中,旨在在降低采样过程中减小信息损失,并在上采样过程中增强细节重建。

在编码器中,作者首先使用DTCWT将特征图分解成高和低频分量,实现下采样,同时减小信息损失。

在解码器中,作者利用iDTCWT从下采样特征中重构更高分辨率特征图。

使用nnU-Net框架在Retina Fluid、Brain Tumor和Liver Tumor分割数据集上的评估表明了所提出Spectral U-Net的优越性。

1 Introduction

深度神经网络(DNNs)已成为医学图像分割的主导方法,该任务为每个图像像素分配一个类别标签。U-Net [16] 是一种卷积神经网络(CNN),在编码阶段通过卷积层生成多尺度的特征图,并在解码阶段逐步恢复各个分辨率的特征图。在编码阶段,下采样用于降低卷积层生成的特征图的空间分辨率,从而减轻计算负担并提高网络的整体效率。下采样聚合周围像素的信息,捕捉不同粒度和尺度的特征。此外,下采样有助于网络降低对图像特征位置的细微变化敏感性,减小特征中的噪声,从而提高提取的特征图的鲁棒性。下采样扩大了更深层神经元的感受野,使神经元能从输入图像的更广区域获取信息,从而使网络能更好地识别复杂的图案和关系。

在解码阶段,上采样用于恢复原始的空间分辨率,使网络能够对像素级进行详细的预测。

作为DNNs的基本过程,已有多种下采样机制被提出,其中最常见的包括最大池化 [24],平均池化 [10],混合池化 [27],随机池化 [28] 和步长卷积 [10]。尽管这些下采样技术简单且有效,但它们可能会引起断点和对细节的损失,特别是在处理医学图像中的小物体(如肿瘤)时,这会阻碍上采样过程中空间分辨率的准确恢复。

为解决上述挑战,作者提出了一种新的框架,称为光谱 U-Net,旨在为医学图像分割提取分级特征图。作者的框架旨在在降低下采样和解样过程中减小信息损失。作者的主要贡献有两个方面。首先,在编码阶段,作者构建了一个Wave-Block来执行下采样。作者首先将输入特征图分解为两个部分,即低频和高频部分,利用对偶树复杂小波变换(DTCWT) [17]。这些部分然后通过共享卷积滤波器处理,结果在降低空间分辨率的同时增加通道数。这种方法使网络在保留池化优势的同时克服信息丢失问题。其次,在解码阶段,作者设计了一个反Wave-Block(iWave-Block)用于分辨率重建。iWave-Block使用反Dual Tree Complex Wavelet Transform(iDTCWT)从下采样的小波系数重构原始输入,同时保留所有相关信息并恢复空间分辨率,从而减小信息丢失。

作者在基于nnU-Net框架 [9] 的多个数据集(如Retina Fluid [3],BRATS 2017 [13] 和 LiTS 2017 [2])上评估作者的新方法,并证明与现有方法相比具有优势。

2 Related Work

Down-sampling and Up-sampling

为降低特征映射的计算效率并提取多尺度特征,采用了降采样方法。卷积池化[24]和平均池化[10]是两种常用的方法,它们在卷积池化区域内选取最大值和平均值,分别为池化区域内的较弱信号丢弃,可能会忽略局部详细信息[27]。平均池化则为池化区域计算平均值,这可能导致池化区域内强激活的权重降低[28]。在[27]中,提出了结合最大池化和平均池化的混合池化操作以提高图像分类准确率[27]。

在[28]中,提出了一种随机池化方案,根据池化区域内的活动以多nomial分布随机选择池化区域内的像素[28]。步长卷积[10]是另一种参数化降采样特征映射的方法。

Wavelet Transform in Deep Learning

小波变换已在时序信号处理和图像预处理与后处理中得到了广泛应用。离散小波变换(DWT)将输入信号分解为具有不同频率和尺度的多个系数。接着,每个这样的系数独立地处理并使用反离散小波变换(iDWT)重建原始信号或特征。在[1]中,它提出利用小波子带来增强图像修复的性能。一种方法被提出,通过利用低分辨率输入的小波子带来恢复高分辨率图像细节[6]。

在[12]中,它给出了一个多级小波变换,用于扩展图像修复任务的空间感受野。小波池化被提出用于进行特征的第二级分解,舍弃第一级子带以减少图像分类中的特征维度[22]。在[5]中,小波变换被用于利用光谱信息进行纹理分类和图像标注。在[11]中,CNNs与小波进行集成,进行噪声鲁棒图像分类。在[26]中,小波变换被用于减少视觉Transformer中键值对的时空分辨率,从而降低计算成本。然而,这些方法分解图像低频和高频成分的能力仍然受到限制,它们捕获细粒度图像细节的能力不足。

在本文中,作者将双树复小波变换(DTCWT)[17]融入CNN框架中,以进行降采样和升采样,目的是减轻传统方法(如最大池化和步长卷积)可能带来的信息损失。具体而言,作者构建一个Wave-Block并沿特征图的每个方向进行DTCWT,从而实现无损空间分辨率降低。在重建阶段,作者引入一个iWave-Block,并利用反DTCWT(iDTCWT)根据前一级的输出重建特征。

3 Method

在这一部分中,作者首先介绍了双树复小波变换(DTCWT)[17]和作者提出的波块和i波块。作者提出的框架的整体结构如图1(a)所示。在编码阶段,作者使用波块逐步降低特征图的空间分辨率并增加通道数。在解码阶段,作者使用i波块逐步重构随后的分割空间分辨率。

picture.image

Dual Tree Complex Wavelet Transform

小波离散变换(DWT)将给定的信号分解为一组子带,每个子带都代表对应频率子带中的系数,描述信号在该子带中的特征。DWT通常存在位移不变性和二、高维缺乏方向性。双树复小波变换(DTCWT),是DWT的扩展,在位移不变性和方向性选择性方面优于DWT。它首先将输入特征X分解为低频和高频成分和:

其中和分别表示低频和高频的小波变换,和表示相应的滤波器,是分解的层数,表示方向选择性。

Wave-Block

在深度神经网络(DNNs)中,最大池(max pooling)和平均池(average pooling)都是不可逆的,可能导致信息损失。为了降低这一问题,作者设计了一个可逆降采样块 ,通过使用频谱分解实现可逆降采样,如图1(b)所示。

对于输入特征图 ,其中 、 和 分别表示 的通道、高度和宽度。DTCWT 首先将 分解为低频分量 和高频分量 ,,,其中 2 表示实部和虚部,6 表示 DTCWT 光谱变换的六个方向:,,,, 和 。然后,作者使用像素换位机制 [18] 将 重新排列为 ,其中 。接下来,作者重新排列 和 :

其中 ,然后用卷积块对其进行处理:

其中 , 表示卷积的参数, 和 分别表示批量归一化和 ReLU 激活。经过一个 Wave-Block,作者降低了 的空间分辨率并翻倍了通道数。然后将 发送到下一个 Wave-Block 进行进一步处理。Wave-Block 的 PyTorch 风格伪代码如下:1:

iWave-Block

在解码阶段,作者逐步将降采样的特征图恢复到原始的空间分辨率,如图1(c)所示。

给定降采样的特征图 , 首先将其重新排列成低频和高频分量 和 , 其中 且 。 然后被重新排列成 :

其中 。随后,作者使用逆 DTCWT 变换重构空间分辨率:

其中 表示逆 DTCWT 变换,。接下来,作者将 和跳跃连接的特征图 (channel dimension) 沿通道维度拼接,并将其馈送到卷积层:

其中 表示 ReLU 激活, 表示批量归一化, 表示卷积的参数,。经过 iWave-Block 处理后,作者提升了图像的空间分辨率并减半了通道数。因此,一个 iWave-Block 可以被视为一个上采样过程。 iWave-Block 的 PyTorch 风格伪代码见代码1。

4 Experiments

Datasets

作者在以下数据集上评估所提出的频谱U-Net:

视网膜液体:作者使用公共的RETOCH和自有的数据集。标签类别包括视网膜内液体(IRF),黄斑囊肿液体(SRF)和色素上皮脱离(PED)。RETOCH数据集来自三个OCT扫描仪:Cirrus,Spectralis和Topcon,分别有24、24和22个受试者。作者的自有数据集包含Spectralis扫描仪的100个OCT受试者。使用的视网膜液体数据集是将RETOCH和自有数据集合并得到的。

BRATS 2017:这是一款用于大脑肿瘤分割的MRI数据集。它包括多种MRI扫描:T1加权,T1对比增强,T2加权和FLAIR图像。目标包括三个肿瘤亚区域:增強肿瘤,周圍水肿和坏死和非增强肿瘤核心。标注被合并在三个嵌套的子区域:全肿瘤(WT),肿瘤核心(TC)和增強肿瘤(ET)。

LiTS 2017:此数据集使用131个对比增强的腹部CT体积用于训练/验证,70个用于测试。标注包括肝部和肿瘤。

Experimental Setup

作者所有的实验都在NVIDIA RTX 6000 Ada GPU上进行。作者基于nnU-Net[9]框架测试所有的模型。具体地,作者使用带有动量参数0.99的梯度下降优化器对模型进行优化。作者采用0.9的幂函数学习率衰减,初始学习率为0.01。作者将每个数据集随机分为训练、验证和测试集,比例为0.7:0.1:0.2。作者报告DSC(Dice相似系数)和HD95(Hausdorff距离(95百分位数))得分。每个实验运行5次,结果的平均值被报告。

Results

表1显示了在视网膜流动(Retina Fluid)数据集上的实验结果。作者与几种典型的方法进行了比较:(1)曲率[23];(2)MsTGANet[20];(3)CPFNet[4];(4)DconnNet[25];(5)Swin UNETR[8];(6)nnU-Net[9]。如图所示,作者的光谱UNet在保留降采样和升采样特征图中的信息方面,分别相较于nnU-Net提高了1.28%和2.35%,其IRF的DSC评分与nnU-Net相当(低0.08%)。这表明作者的方法在降采样和升采样过程中,对信息的保留非常有效。

picture.image

表2展示了在BRATS 2017数据集上的比较结果,与几种典型的方法进行了对比:(1)TransBTS[21];(2)SegResNet[14];(3)VT-UNet[15];(4)Swin UNETR[8];(5)nnU-Net[9]。作者发现ET和WT的DSC评分分别优于Swin UNETR[8]的0.32%和0.78%。与Swin UNETR相对于nnU-Net的改进(即0.11%和-0.04%)相比,作者改进的成效证明了作者光谱UNet在降采样和升采样过程中,信息的保留能力十分强大。

picture.image

表3给出了在LiTS 2017数据集上的比较结果,与几种典型的方法进行了对比:(1)TransBTS[21];(2)KiU-Net[19];(3)Swin UNETR[8];(4)nnU-Net[9]。作者的光谱UNet在肝癌症和肝的DSC评分上,分别优于第二好的Swin UNETR的0.7%和0.45%。与Swin UNETR相对于nnU-Net的改进(即0.23%和0.34%)相比,作者改进的成效强调了作者的方法在降采样/升采样特征图过程中的信息损失减少能力。

picture.image

Ablation Study

作者在视网膜液体数据集上进行了消融实验,以评估作者的方法的有效性。IRF、SRF和PED的DSC评分报告在表4中。ConvBlock指的是传统的最大池化卷积块,Linear-I表示线性插值。作者观察到,作者提出的Wave-Block与ConvBlock相比,始终提高了DSC评分。请注意,仅使用ConvBlock和iWave-Block一起时的改进有限,因为在最大池化过程中信息已经丢失,很难在向上采样过程中恢复。这进一步证明了在降采样过程中使用作者的DTCWT减少信息损失的有效性。此外,作者将DTCWT与常用的Haar小波变换[7]进行了比较。实验结果表明,与在降采样过程中更有效分解高频和低频分量相比,DTCWT在信息保留方面更具优势。

picture.image

Qualitative Results

部分来自Retina Fluid数据集的质性示例如图2所示。这些示例证明作者的方法能够捕获小物体和边缘区域(用虚线红框标出)的精细细节,而这些细节被nnUNet、Swin UNETR和DconnNet所遗漏。

picture.image

Parameters and Computation Costs

要检验作者方法的成本,作者在表5中报告了参数、浮点运算(FLOPs)和推理时间(执行一次卷积神经网络预测所需的时间,以秒为单位)。从表中可以看出,作者的参数、FLOPs和推理时间与nnU-Net相似,这表明了作者方法的高效计算能力。

5 Conclusions

本文提出了一种名为Spectral U-Net的新颖分割网络,该网络利用双树复小波变换(DTCWT)进行信息无损降采样,利用正向双树小波变换(iDTCWT)进行上采样过程。

作者在降采样阶段利用DTCWT将特征图分解为低频和高频分量,在解码阶段利用iDTCWT进行重建,从而减少了传统池化的信息损失。

在Retina Fluid、BRATS 2017和LiTS 2017分割数据集上的实验结果表明,新方法能有效捕获被nnUNet、Swin UNETR和DconnNet 遗漏的微小物体和边缘细节,同时保持计算效率。

参考

[1].Spectral U-Net: Enhancing Medical Image Segmentation via Spectral Decomposition.

点击上方卡片,关注 「AI视界引擎」 公众号

0
0
0
0
关于作者
关于作者

文章

0

获赞

0

收藏

0

相关资源
CV 技术在视频创作中的应用
本次演讲将介绍在拍摄、编辑等场景,我们如何利用 AI 技术赋能创作者;以及基于这些场景,字节跳动积累的领先技术能力。
相关产品
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论