新加坡国立大学提出 AIC-UNet | 解决图像精准分割问题再上一层 ！

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

强制实施关键解剖学特征，如器官的数量、形状、大小以及相对位置，对于构建一个健壮的多器官分割模型至关重要。

当前尝试合并解剖学特征的方法包括使用资源和数据密集型的模块如自注意力来扩大有效的感受野（ERF）大小，或者引入特定器官的拓扑规则器，但这些方法可能无法扩展到多器官分割问题，其中器官间关系也起着巨大的作用。

作者提出了一种新方法，通过用可学习的解剖先验条件模型预测，在任何现有的编码器-解码器分割模型上施加解剖学约束。

更具体地说，对于给定的腹部扫描，编码器的一部分使用薄板样条（TPS）网格插值在空间上扭曲一个可学习的先验，使其与给定的输入扫描对齐。在解码阶段，将扭曲的先验整合进来，指导模型进行更具解剖学信息的预测。

代码可在https://anonymous.4open.science/r/AIC-UNet-7048获取。

1 Introduction

遇到报告称其性能与放射科医生相当甚至超越的人工智能模型越来越普遍。

然而，人工智能模型在不久的将来取代放射科医生的可能性非常小。尽管在大多数情况下AI模型表现良好，但它们仍可能做出放射科医生永远不会做出的解剖学上错误的预测。例如，AI可能预测食道，即从喉咙输送食物和液体至胃的肌肉管，是分离的。或者，AI可能会错误地将胫骨识别为股骨，因为这两块骨头在局部层面上看起来可能相似。

这些例子说明当前的人工智能模型在学习关键解剖特征方面并不可靠。

那么，即使在训练过程中接触到数十万个实例之后，是什么阻止了当前的人工智能模型识别关键的解剖特征呢？目前，完全自主的人工智能分割模型仅根据输入扫描来检测器官。

因此，作者期望AI能够以完全数据驱动的方式自行捕捉这些解剖限制。然而，这些解剖特征不仅是全局特征，比起局部特征来说学习难度大得多，而且模型甚至可能并不将它们视为严格的限制。因此，对当前人工智能模型来说，学习或融合解剖特征是一个巨大的挑战。

已经提出了几种方法来融入或更好地学习解剖学约束。这些方法可以广泛分为以下两类：1）扩大有效的感受野（ERF）或2）施加拓扑约束。许多研究考虑了图神经网络（GNN）[15; 9]和自注意力网络[17; 7]，这些网络比标准的卷积神经网络（CNNs）更适合于发现全局特征。然而，这些网络通常需要更多的训练数据以获得更好的泛化能力[6]，这在医学领域可能是一个具有挑战性的要求。一些研究考虑将分割重新定义为网格变形任务[1; 10]，通过学习变形模板网格自然提供更平滑的轮廓预测。然而，这种方法在表示复杂结构时存在困难。拓扑正则化技术[14; 8]，虽然在特定解剖学挑战中有效，但当任务扩展到多器官分割时，限制了技术的泛化能力，这时器官的相对位置扮演着更重要的角色。

作者提出的模型，即解剖学信息级联UNet（1.），旨在融合全局解剖学先验信息，同时不对标准的编码器-解码器分割网络进行重大修改。

该模型融合了解剖学特征，但并不依赖于资源密集型的全局上下文学习器，或者不能处理器官间关系的拓扑正则化器。相反，作者引入了一个可学习的额外参数，称为“先验”，它可以进行空间变形以匹配患者的解剖结构。变形的先验在预测过程中起到软约束的作用。更具体地说，图像编码器输出的一部分被用作薄板样条（TPS）网格变形的控制点。TPS将可学习的先验变形以匹配患者的解剖结构。在解码阶段，变形的先验被整合进来，引导解码器进行更加符合解剖信息的预测。为了进一步提高复杂目标的变形精度，作者用裁剪后的局部图像块重复相同的过程，从而形成一个全局-局部级联网络。

本文的贡献总结如下：

1. 作者 Proposal 通过用变形的解剖先验条件来调整其预测，以提高多器官分2. 割模型的鲁棒性。

2. 作者提出了一种全局-局部级联变形方法，用以提高复杂物体的变形精度。

作者提出了一种激活最大化技术，以学习一个通用先验，而不是使用固定的解剖模板。

2 Prior Works

现有用于增强解剖学特征学习的方法侧重于扩大有效感受野（ERF），将分割重新定义为网格变形，或通过正则化器施加拓扑约束，每种方法都有其自身的局限性。

2.0.1 Broadening ERF

GNNs [15] 和自注意力网络 [17] 是可以获取比标准CNN更大ERF的模型。因此，这些模型更适合学习数据中的远距离依赖性，使它们成为学习解剖特征的理想选择 [3, 7, 13]。然而，在实际应用中，由于可用于监督长距离依赖性学习的数据有限，这些模型可能难以有效地学习解剖先验知识。

2.0.2 Mesh Deformation

网格变形[10, 1]天然地提供了比传统像素预测更平滑的轮廓预测，但在表现复杂结构时遇到了困难。一个潜在的解决方案是将基于网格的分割与基于像素的方法相结合。然而，由于这两种方法在目标表示上的本质差异，这种方法也带来了挑战。

2.0.3 Topology regularization

不同器官的形状可以通过它们的拓扑特征来描述，例如对于三维体积来说，物体的数量和空腔的数量。有些研究使用拓扑约束来规范网络预测。基于拓扑的技术通常针对特定的解剖学挑战进行定制，这降低了它们的泛用性。此外，这些方法倾向于优先惩罚局部概念，比如孔洞的数量或局部的连通性，而没有解决更广泛的考虑，如器官的形状及其之间的关系。

3 Method

3.0.1 Network Overview

在核心层面，AIC-UNet（如图1所示）是一个级联网络，它需要全局视图和局部视图来生成一个全面的地方多器官预测，其中表示器官的数量。

最初，模型接收一个缩放后的全局视图 作为输入，生成初始的粗略全局预测 。除了分割模型的标准编码器-解码器结构外，AIC-UNet 还融合了一个自由参数 以及3个计算模块：PriorEncoder, Deform 和 。 被优化以表示一般的解剖结构。通过一个形变模块 Deform 对 进行空间形变，生成 ，该值紧密跟随 的解剖结构。形变程度由从 Encoder 和 PriorEncoder 提取的并接特征进行调节。此过程在模型的局部片段中重复，采用局部视图 和局部先验 ，以及作为输入先验的裁剪和重缩放的全局预测 。局部模型用于细化形变的全局先验，产生形变的局部先验 。

3.2.2 Deformation block

图1：AIC-UNet概览。AIC-UNet是一个级联网络，结合了全局和局部视图，用于全面的多器官预测。初始输入 产生粗糙的全局预测 ，通过 进行增强，这是通过Deform从可学习的先验 得到的空间变形解剖结构。此过程在模型的局部段中重复以进一步增强，采用局部视图 和局部先验 。

图2：变形块将可学习的先验解剖结构 变形为特定患者的解剖结构 。

如图2所示，Deform块接收两个输入：估计的源控制点 和先验 ，生成一个变形的先验 。作者使用TPS变形[2]，它允许使用稀疏控制点进行非线性变形，来变形通用解剖结构 。

TPS变形按以下方式工作：给定一组预定义的密集目标控制点序列 和稀疏源控制点 ，$N<<m$，通过最小弯曲能量匹配控制点 $\mathbf{p}^{(i)}\_{target}\mapsto\mathbf{p}^{(i)}\_{source}$="" 确定一个唯一的网格变形函数="" $\mathcal{d}$。控制点的="" $h$、$w$="" 和="" $d$="" 坐标的匹配，连同正则化条件，产生了三组系数="" $(a^{(1)},\cdots,a^{(n+4)})$，$(b^{(1)},\cdots,b^{(n+4)})$="" $(c^{(1)},\cdots,c^{(n+4)})$，使得一般的目标点="" $\mathbf{p}="(h,w,d)$" 映射到="" $(\mathcal{d}\_{h}(\mathbf{p}),\mathcal{d}\_{w}(\mathbf{p}),\mathcal{d}\_{d}(="" \mathbf{p}))$。<="" p="">

(1b)

在这里， 是一个核函数。关于优化TPS系数的更多细节在附录中给出。

3.2.2 Learnable Prior

拥有一个真实的器官解剖结构作为全局先验，极大地提高了随后变形的全局和局部先验的准确性。与其他基于图谱的分割方法[10; 1]不同，后者从训练集中分配一个任意的真实解剖结构或简单的结构，如球体，AIC-UNet在训练过程中学会寻找最优的全局先验。这是通过将全局先验 做为一个自由参数来实现，其大小为 ，这匹配了全局视图的空间维度。作者在通道维度 上应用Softmax，以限制范围在 之间。

作者进一步解释了在“损失函数与优化”小节中用于加速先验学习的优化技巧。

3.2.3 Aggregation of Prior

变形先验 通过Squeeze Excitation（SE）注意力模块与解码器块的输出结合。对注意力调节后的特征应用卷积，以匹配后续解码器块所需的通道尺寸。

3.3.2 Loss Function and Optimization

AIC-UNet被训练以最小化Dice损失，它是全局和局部 Level 上计算的Soft-Dice和交叉熵损失的组合，以及用于稳定源控制点估计的正则化项。损失定义如下：

在这里， 和 分别代表全局和局部的真实解剖结构。 和 是控制 正则化程度的系数。

直接优化全局先验 以及其他模型参数会导致收敛速度变慢。作者假设这是由于先验与TPS控制点之间的相关性造成的。例如，如果预测的解剖结构小于真实值，可以通过两种方式减少误差：1) 在TPS变形中缩小源控制点；2) 放大全局先验。这种相关性可能会混淆优化优先级。作者通过交替优化模型参数和全局先验来防止混淆。

4 Dataset and Experimental Setup

4.0.1 Dataset

作者使用了公开可获取的全腹器官数据集（Whole abdominal ORgan Dataset，简称WORD）[12]。WORD由150个匿名化的CT扫描组成。每个扫描包含159-330个切片，每个切片的分辨率为像素，平面内分辨率为0.976 mm 0.976 mm。切片间隔在2.5到3.0毫米之间变化。WORD提供了16个器官的标注，包括肝脏、脾脏、肾脏以及各种消化器官。为了训练、验证和测试的目的，该数据集被随机划分为100个扫描用于训练，20个扫描用于验证，以及30个扫描用于测试。

4.0.2 Experimental Setup

像素强度被截断在之间，并使间隔标准化为。轴向方向（维）通过零填充以具有相同的体积大小。全局视图和全局 Mask 通过因子进行下采样。局部视图的维度设置为。作者使用AdamW [11]优化器，并采用线性 Warm up 余弦退火。最大学习率和权重衰减分别设置为3e4和1e5。对于先验优化的学习率，设置为1e3。每500次迭代，作者对先验进行100次迭代的训练。模型训练350个周期，每个周期包含200次迭代。批量大小设置为2。在损失（2）中对都使用1e8以规范TPS控制点学习。

4.0.3 Baseline Methods

作者与两种 Baseline 方法进行了比较。第一种是标准的UNet，它构成了作者局部分割网络的基础框架。第二种是级联UNet（CUNet），其结构与作者的AIC-UNet相同，不同之处在于它没有共同的先验知识 。对于CUNet，作者在全局分割网络上使用自注意力机制，并从全局网络中的裁剪预测中关注来自局部分割网络的特征。

5 Experimental Results

5.0.1 Segmentation Metrics

作者通过三种指标来衡量分割性能：骰子分数（Dice）、归一化表面骰子（NSD）[16] 和 95% 的豪斯多夫距离。作者的 AIC-UNet 在这三个指标上均取得了最佳的平均结果。

5.0.2 Visualization of Deformed Prior

图3展示了通过TPS变形块学习的全局先验解剖结构和患者特异性的变形解剖。该图说明所学习的全局先验与作者对通用器官解剖的理解非常吻合。此外，先验解剖也成功地变形为不同的患者特异性解剖。例如，图中位主体的变形解剖准确地表示了该主体相对较短的躯干和较小的腰臀比。

5.0.3 Qualitative Comparison

表1：在WORD数据集上提出模型与基准模型的比较。

作者的AIC-UNet中学习到的公共先验上的注意力机制可以促进解剖学上的准确分割。这一点得到了图4中的结果支持。在图4a-4b中，UNet 错误地将靠近膝盖的骨头分割成了股骨头，而AIC-UNet提供了更准确的分割。在胆囊分割的情况下，如图4c-4d所示，UNet的预测多出了一个组成部分，并且与真实位置（以透明红色标出）有所偏差，而AIC-UNet正确识别了组成部分的数量及其大致位置。## 6 结论与未来研究

图4：AIC-UNet与UNet在股骨头（黄色和蓝色）和胆囊分割上的定性比较。透明叠加了肝脏和胆囊的真实值。

作者提出了AIC-UNet，这是一个编解码分割模型，通过使用先验变形来利用解剖学信息。为了改进这个模型，未来的研究可以关注两个领域。首先，通过开发更有效的目标控制点选择策略来提高TPS变形性能。其次，通过设计一个更强大的特征聚合模块，能够将从先验变形整合到解码器块中的信息。## 附录

Solution of Thin Plate Spline Coefficients

鉴于已估计的源控制点 ，3D-TPS 将一般目标点 映射到 。

(3b)

在这里， 是一个核函数，，，以及 是需要优化的系数。下面作者将解释这三组系数是如何计算的。作者以 -坐标系数为例，而 和 坐标系数的计算方式也类似。

薄板函数（3a）需要计算 个系数。尽管该函数与核函数 的高度非线性相关，但相对于系数而言，该函数是线性的。因此，这些系数具有闭式解。

令向量 ，其中 是第 个源控制点的 坐标。同时，定义矩阵。

在公式中，，，，以及分别代表源控制点的-，-，和-坐标，而是一个大小为的零矩阵。那么系数由以下给出：

参考

[1].AIC-UNet: Anatomy-informed Cascaded UNet for Robust Multi-Organ Segmentation.

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