[论文解读] PC-SwinMorph: Patch Representation for Unsupervised Medical Image Registration and Segmentation

Lihao Liu, Zhening Huang+3 more|arXiv (Cornell University)|Mar 10, 2022

Advanced Neural Network Applications被引用 13

一句话总结

PC-SwinMorph 提出了一种统一的无监督框架，用于医学图像配准与分割，采用两种新颖的基于图像块的策略：一种用于增强局部与全局特征表示的图像块对比学习方法，以及一种3D窗口/移位窗口多头自注意力模块，以消除图像块分割带来的伪影。该方法实现了最先进性能，在 CANDI 和 LPBA40 数据集上相比基线模型，Dice 分数最高提升 4%。

ABSTRACT

Medical image registration and segmentation are critical tasks for several clinical procedures. Manual realisation of those tasks is time-consuming and the quality is highly dependent on the level of expertise of the physician. To mitigate that laborious task, automatic tools have been developed where the majority of solutions are supervised techniques. However, in medical domain, the strong assumption of having a well-representative ground truth is far from being realistic. To overcome this challenge, unsupervised techniques have been investigated. However, they are still limited in performance and they fail to produce plausible results. In this work, we propose a novel unified unsupervised framework for image registration and segmentation that we called PC-SwinMorph. The core of our framework is two patch-based strategies, where we demonstrate that patch representation is key for performance gain. We first introduce a patch-based contrastive strategy that enforces locality conditions and richer feature representation. Secondly, we utilise a 3D window/shifted-window multi-head self-attention module as a patch stitching strategy to eliminate artifacts from the patch splitting. We demonstrate, through a set of numerical and visual results, that our technique outperforms current state-of-the-art unsupervised techniques.

研究动机与目标

• 为解决无监督医学图像配准与分割的局限性，这些方法常因缺乏监督而难以捕捉精细解剖细节。
• 通过利用基于图像块的空间结构，提升无监督学习中的特征表示能力，实现更优的局部与全局解剖建模。
• 通过一种新颖的图像块拼接机制，结合 3D W-MSA 和 SW-MSA 模块，减轻 3D 医学图像中基于图像块表示带来的伪影。
• 在无需真实形变场或分割掩码的情况下，同时实现配准与分割任务的最先进性能。
• 证明基于图像块的表示相较于全局标记嵌入，在复杂医学图像任务中更具意义。

提出的方法

• 该框架采用 VoxelMorph 作为主干网络，将其编码器替换为基于图像块的表示策略，以增强特征学习能力。
• 引入一种图像块级别的对比学习策略，通过跨图像比较图像块，强制实现局部一致性并获得更丰富的特征表示。
• 在编码器/解码器之后应用 3D 窗口/移位窗口多头自注意力模块（3D W-MSA 和 3D SW-MSA），以拼接图像块并恢复长距离依赖关系。
• 将对比损失与图像块拼接机制整合进统一框架，同步提升配准与分割性能。
• 3D SW-MSA 模块作为主干网络后的单层应用，最大限度减少计算成本，同时实现最优的特征融合。
• 模型通过无监督损失函数端到端训练，无需真实形变场或分割掩码。

实验结果

研究问题

• RQ1 与基于全局标记的方法相比，基于图像块的表示是否能显著提升无监督医学图像配准与分割性能？
• RQ2 引入图像块级别的对比学习策略是否能增强特征表示并改善细粒度解剖对齐？
• RQ3 3D 窗口/移位窗口多头自注意力机制是否能有效减少图像块分割带来的伪影并提升跨图像块的特征交换？
• RQ4 将图像块对比学习与图像块拼接相结合，能在真实医学数据集上多大程度上提升配准与分割性能？
• RQ5 所提出的方法在实现无监督设置下的最先进性能的同时，是否具有计算效率？

主要发现

• 所提出的 PC-SwinMorph 框架在 CANDI 和 LPBA40 数据集上，相比基线模型，Dice 分数最高提升 4%。
• 仅使用图像块对比学习策略，Dice 分数相比对比学习基线提升 2%，证明其在特征增强方面的有效性。
• 在 CANDI 数据集上，3D W-MSA 和 3D SW-MSA 模块的集成使配准性能提升 3.06%，分割性能提升 2.65%。
• 该模型在 CANDI 数据集上实现了 0.817 的分割 Dice 分数与 0.812 的配准 Dice 分数，优于所有最先进无监督方法。
• 计算开销增加极小，仅增加 0.6 GFlops（从 416.04 增至 416.10 GFlops），证实了方法的高效性。
• 视觉与数值结果均表明，基于图像块的方法在捕捉精细解剖细节方面，比全局表示方法更具鲁棒性。