Eformer: Edge Enhancement based Transformer for Medical ImageDenoising--用于医学图像去噪的基于边缘增强的 Transformer

作者在摘要中提到,Eformer使用了用于医学图像去噪的 Transformer block构建Encoder-Decoder Net。

Transformer block中使用了基于非重叠窗口的自注意力,这可以降低计算要求;

结合了可学习的 Sobel-Feldman 算子用以增强图像边缘,并配套了有效的方法将它们连接到架构的中间层。

在AAPM数据集上实现了43.487 PSNR、0.0067 RMSE 和 0.9861 SSIM的成绩。

介绍

作者提到他们的想法受益于ViT的 patch embedding(块编码)操作,随即提到本论文和核心点:

        1. 将可学习的 Sobel 滤波器用于边缘增强,从而提高了整体架构的性能。

        2. 按照残差学习范式对网络进行了广泛的训练。为了证明残差学习在图像去噪任务中的有效性,作者还使用定性方法展示了结果,其中我们的模型直接预测去噪图像。在医学图像去噪中,残差学习明显优于传统学习方法,在传统学习方法中,直接预测去噪图像变得类似于制定恒等映射。

相关工作

        作者在此列举前人贡献,为老生常谈,但有一点值得注意:作者的idea受到了一种去噪Transformer -- Uformer的启发,它在前馈网络中采用了基于非重叠窗口的自注意力和深度卷积来有效的捕获局部的上下文信息。

而作者自述的贡献是:他们用了一种新颖的方式集成了边缘增强模块和类似于Uformer的架构。

方法

Eformer: Edge Enhancement based Transformer for Medical ImageDenoising--用于医学图像去噪的基于边缘增强的 Transformer_第1张图片 图1 架构详细描述

3.1. Sobel-Feldman Operator:

       

Eformer: Edge Enhancement based Transformer for Medical ImageDenoising--用于医学图像去噪的基于边缘增强的 Transformer_第2张图片 图2. 实现中有四组不同的 Sobel 滤波器,其中α是可学习参数

 

Eformer: Edge Enhancement based Transformer for Medical ImageDenoising--用于医学图像去噪的基于边缘增强的 Transformer_第3张图片 图 3. 使用 Sobel 滤波器对图像进行卷积后获得的结果示例。输入(左)和边缘增强图像(右)。

 

        作者使用Sobel-Feldman 算子,也称Sobel滤波器,是一种专门用于边缘检测的算法,它原本有两个版本,分别是:垂直和水平,作者在他的架构中加入了拓展的对角线版本算子,检测效果如图3所示。包含边缘信息的图像特征图集与输入影像和网络的其他部分有效连接,具体参见图1.

3.2. Transformer based Encoder-Decoder

        该部分的灵感主要来自于Uformer。在每个编码和解码的阶段,卷积特征图都通过一个局部增强窗口(LeWin)Transformer block。这个block包括一个基于非重叠窗口的多头自注意力模块(W-MSA)和一个局部增强馈送前向网络(LeFF)。总体的数学表达公式如下,LN表示Layer Normlization:

         由图1所示,在每个编码阶段,本transformer block在 LC2D 块之前被应用,在每个解码阶段在 LC2U 块之后被应用,也用作瓶颈层。

3.3. Downsampling & Upsampling

        对于下采样,作者先否决了常规的池化此层,转而选择了加大stride(步长)的卷积层用于下采样,更具体地说,他使用 3 × 3 的内核大小,步幅为 2,填充为 1。

        对于上采样,作者选择了转置卷积,它可以重建空间维度并学习自己的参数,但是它也存在可能造成棋盘伪影的问题。作者参考资料认为避免不均匀重叠可以缓解这一问题,于此确定了参数:使用 4 × 4 的内核大小和 2 的步幅。

3.4. Residual Learning

        没什么好说的,就是输出残差图会更容易优化。

3.5. Optimization

        本文使用了多个损失函数,最初是MSELoss:

         但是,它往往会产生不必要的伪影,例如过度平滑和图像模糊。为了克服这个问题,我们同时加入了基于 ResNet [12] 的多尺度感知 (MSP)Loss:

        ResNet-50 主干被用作特征提取器 φ。具体来说,删除了在 ImageNet 数据集上预训练的 ResNet-50 的池化层,保留了权重 (θ) 被冻结之后的卷积块。为了计算感知损失,将模型输出R(xi)经过计算 xi - R(xi) 后传递给提取器。在此之后,从主干的四个阶段提取特征图,如Ednn中所做的那样(之前提到要去掉ResNet-50的池化层,提取特征图的位置也就是这原本池化层的四个位置,下图4是从Ednn中扣到的图)。

Eformer: Edge Enhancement based Transformer for Medical ImageDenoising--用于医学图像去噪的基于边缘增强的 Transformer_第4张图片 图4. 多尺度感知损失。它基于 ResNet 50,包含 4 个主要阶段。有关此模型的更多详细信息,请参阅论文

         到此得到最终的损失函数:

 3.6. Overall Network Architecture

        输出LDCT图像首先通过Sobel 滤波器以产生S(I),然后是GeLU激活函数。

        作为编码阶段的一部分,在每个阶段,我们将输入通过 LeWin Transformer Block,然后通过与 S(I) 的连接和随后的卷积操作;然后使用大步长的卷积进行下采样;编码后,在瓶颈处,我们将编码后的特征图传递给另一个 LeWin Transformer Block。

        在解码器的每个阶段,通过反卷积后,较早的下采样 S(I) 与上采样的特征图连接,然后通过卷积块。解码器阶段可以看作是编码器阶段的对立面,具有共享的 S(I)。

        解码后生成的最终特征图然后通过“输出投影”块以产生所需的残差。这个“输出投影”是一个卷积层,它简单地将 C 通道特征图投影到 1 通道灰度图像。

        在我们的实验中,我们将 LeWin 块的深度、注意力头和编码器-解码器阶段的数量设置为 2。

附文

Eformer: Edge Enhancement based Transformer for Medical ImageDenoising--用于医学图像去噪的基于边缘增强的 Transformer_第5张图片

Parameter Details and Network Training

        为了训练,我们使用 Pydicom 库将图像像素数据从 Dicom 文件提取到 NumPy 数组,然后将 NumPy 数组中的像素数据从 0 缩放到 1,以避免不同 CT 扫描的像素数据的异质跨越。        

        该模型使用 ADAM 优化器使用 0.00002 的学习率和默认参数进行训练。该模型使用 128 × 128 像素的输入大小进行训练,选取方法是从其原始大小 512 × 512 像素调整每个图像的大小。得到的实验结果如图5:

Eformer: Edge Enhancement based Transformer for Medical ImageDenoising--用于医学图像去噪的基于边缘增强的 Transformer_第6张图片 图5. 结果

LeWin Transformer

        LeWin transformer block 如下图所示,包含两个核心设计:

Eformer: Edge Enhancement based Transformer for Medical ImageDenoising--用于医学图像去噪的基于边缘增强的 Transformer_第7张图片

        首先,基于非重叠窗口的多头自注意力(W-MSA),它适用于低分辨率特征图,足以学习远程依赖。其次是局部增强的前馈网络(LeFF),它将卷积算子与传统的前馈网络集成在一起,对于学习局部上下文至关重要。

        在 LeFF 中,图像块首先通过线性投影层,然后是 3×3 深度卷积层。此外,补丁特征被展平并最终传递到另一个线性层以匹配输入通道的尺寸。 对应的方程式如下:

Eformer: Edge Enhancement based Transformer for Medical ImageDenoising--用于医学图像去噪的基于边缘增强的 Transformer_第8张图片

        在 W-MSA 模块中,给定的 2D 特征图 X ∈ C×H×W 被分成 N 个非重叠窗口,窗口大小为 M×M。然后,对每个窗口 X^{i}∈ M^2×C 的扁平特征执行自我注意。假设多头注意力的头部编号为 j,头部维度为 d_{j} = C/j。然后,随之而来的计算是:

Eformer: Edge Enhancement based Transformer for Medical ImageDenoising--用于医学图像去噪的基于边缘增强的 Transformer_第9张图片

         表示第 j 个头的输出。现在,可以将所有头的输出连接起来,然后进行线性投影以获得最终结果。

 笔者读到这里有一个问题,在论文开篇提到的他从ViT的简单标记化中获得的启发到底是什么,希望各位读者可以在评论区指导。

你可能感兴趣的:(transformer,图像处理,计算机视觉)