【技术追踪】UNest：一种用于非配对医学图像合成的新框架（MICCAI-2024）

news2024/10/5 15:32:28

前天看了一篇文章图像分割用diffusion，今天看了篇文章图像合成不用diffusion，你说说这~

传送门：【技术追踪】SDSeg：医学图像的 Stable Diffusion 分割（MICCAI-2024）

UNest：UNet结构的Transformer，一种用于非配对医学图像合成的新框架，涵盖三种模态(MR、CT和PET)，在六项医学图像合成任务中将最近的方法改进了19.30%。

论文：Structural Attention: Rethinking Transformer for Unpaired Medical Image Synthesis
代码：https://github.com/HieuPhan33/MICCAI2024-UNest （即将开源）

0、摘要

非配对医学图像合成的目的是为准确的临床诊断提供补充信息，并解决获得对齐的多模态医学扫描的挑战。
由于Transformer能够捕获长期依赖关系，他们在图像转换任务中表现非常出色，但只是在监督训练中有效，在非配对图像转换中性能下降，特别是在合成结构细节方面。
本文的经验证明，在缺乏成对数据和强归纳偏差的情况下，Transformer会收敛到非最优解。为了解决这个问题，本文引入了UNet结构Transformer（UNet Structured Transformer，UNest）—— 一种新的架构，它包含了结构归纳偏差，用于非配对的医学图像合成。
本文利用SAM模型来精确地提取前景结构，并在主要解剖结构中实施结构注意。这会指导模型学习关键的解剖区域，从而在缺乏监督的非配对训练中改进结构合成。
在两个公共数据集上进行评估，涵盖三种模态(MR、CT和PET)，在六项医学图像合成任务中将最近的方法改进了19.30%。

1、引言

1.1、图像合成的意义

医学影像具有多个模态，不同模态可提供互补的信息，但多次扫描可能是耗时、昂贵的，且有辐射暴露的风险，医学图像合成是一种新思路。（格局打开~）

1.2、现有合成方法局限

（1）大多数合成方法基于有监督的Pix2Pix方法，需要成对数据，不好获取；
（2）CycleGAN是非配对图像转换的开创性工作；
（3）以往的方法采用带有局部归纳偏置的卷积算子，指导模型提取局部特征。这限制了它们捕捉远程空间上下文的能力；

1.3、ViT方法的不足

（1）ViT可建模全局依赖，在分割、超分任务上表现优异；然而，ViT模型在应用于未配对的医学图像合成时很困难；
（2）由于没有归纳偏置，ViT的样本效率较低，在低数据条件下无法注意到鉴别特征；

（a）目前的ViT方法无法在鼻腔内合成复杂的解剖结构；
（b）Transformer方法倾向于关注不太相关的背景特征；
在这里插入图片描述

1.4、本文贡献

（1）本文发现，加入结构引导偏差使Transformer能够专注于鉴别区域，从而增强了非配对图像合成中解剖结构的合成；
（2）提出UNest框架，应用了双重注意策略：前景的结构注意和背景的局部注意；
（3）对MR、CT和PET三种模态的六种图像转换任务进行评估，UNest显著提高了各种解剖结构的准确性；

2、方法

UNest整体框架图：
在这里插入图片描述

2.1、CycleGAN概述

基于CycleGAN，UNest有两个生成器： ${G_{XY}}$ 和 ${G_{YX}}$ ，学习 ${X}$ 和 ${Y}$ 两个域之间的前向和向后映射。

${G_{XY}}$ 和 ${G_{YX}}$ 被训练来欺骗鉴别器 ${D_{Y}}$ 和 ${D_{X}}$ ，训练损失为对抗损失：
在这里插入图片描述
针对未配对的训练，CycleGAN施加了循环一致性损失：

在这里插入图片描述

整体上还是CycleGAN模式：

在这里插入图片描述

2.2、非配对图像合成中的Transformer模型分析

常规的 self-attention 是在整个图中做的：
在这里插入图片描述
本文考虑CNN的局部归纳偏差，从查询（Q） tokens 周围的 ${m×m}$ 窗口中聚合 tokens：

在这里插入图片描述

ResViT和UNETR采用的全局注意使髋关节结构变形，而Swin UNETR采用的局部注意产生伪影：
在这里插入图片描述

2.3、UNet Structural Transformer

与之前的全局注意或局部注意不同，本文的结构注意聚集在主要解剖结构中，在划分patch之后，采用轻量级分类器实现对patch的分类标注。
怎么训练patch分类器呢，使用SAM提取原图的分割结果，取前景最大部分为mask，计算与真实标签的BCE损失，优化patch分类器：
在这里插入图片描述
最终损失为三个损失的加权：

UNest由Structural Transformer（ST）块和一个具有跳跃连接的卷积解码器组成：
在这里插入图片描述

2.4、Structural Transformer模块细节

为了在非配对训练下引导 Transformer，本文采用双重注意策略。对于前景，利用结构注意来学习解剖区域内的关系。对于背景，进行局部关注，实现前景和背景特征之间的有效信息交换。（具体实施还是后面看代码比较直观~）

分别对前景 tokens 和背景 tokens 实施双重注意策略：
在这里插入图片描述

3、实验与结果

3.1、数据集与实施细节

（1）MRXFDG数据集：MR-to-CT、MR-to-PET，37例，224×224；
（2）AutoPET数据集：PET-to-CT，310例，256×256；
（3）两数据集划分：8：1：1；
（4）显卡： 2块 NVIDIA RTX 3090 GPUs ；
（5）优化器：Adam；
（6）epoch：100；
（7）学习率：0.0001，在最后50个epoch线性衰减到0；