论文精读(九)ImplicitAtlas:学习医学成像中的可变形形状模板

CVPR 2022 Open Access Repository

图 3.ImplicitAtlas 概述。 (a) 模型管道。 该模型由隐式模板网络(T)和隐式变形网络(D)组成。 给定一个潜在代码 h,它通过 STE-Softmax 选择一个潜在模板 t,以使用 T 生成模板。它还产生一个潜在变形 d,以使用 D 从模板生成变形场。它们组合起来产生一个占用场, 这是最终的输出。 (b) 解码器的网络架构。 给定一个潜在特征,它以卷积方式构建多尺度特征图。 对于查询坐标 p = (x, y, z),它通过在多尺度特征图上插值来聚合局部和全局特征。 最后,将坐标和插值特征输入 MLP 以进行最终输出。

Abstract

深度隐式形状模型在整个计算机视觉社区中已经很流行,但在生物医学应用中却不太流行。 部分原因是大型训练数据库不存在,部分原因是生物医学注释通常很嘈杂。 在本文中,我们表明,通过在深度学习管道中引入模板,我们可以克服这些问题。 所提出的框架名为 ImplicitAtlas,将形状表示为来自学习模板场的变形场,其中可以集成多个模板以提高形状表示能力,而计算成本可以忽略不计。 对三个医学形状数据集的大量实验证明了其相对于当前隐式表示方法的优越性。

1. Introduction

形状建模是医学图像分析的核心,许多不同的表面表示方法已用于此目的 [16, 36]。 近年来,深度隐式曲面 [8,38,44] 已成为更成熟方法的强大替代方案。 这在整个计算机视觉领域尤其如此,但在生物医学成像子领域则不然。

这归因于生物医学数据集 [27, 56] 带来的具体挑战:对于制造的物体,有大量数据集 [6, 31] 可用于形状模型的训练目的。 对于医学成像中的许多感兴趣的对象(例如器官和病变的形状)来说,这些并不存在。 即使医学图像和相应的 3D 模型可用,由于精确注释的复杂性和费用,模型的质量也低得多。 此外,在单个数据集中,空间分辨率很少是恒定的,并且通常是各向异性的。 人为错误可能会导致标记噪声,并且由于成像过程的限制,器官边界经常被裁剪。 图 2 说明了其中一些困难。

为了解决这些问题,我们提出了 ImplicitAtlas,一种用于医学成像的数据高效隐式形状模型。 在训练过程中,离散网格上表示的样本被作为输入,这是生物医学成像中最常见的表示,并且模型输出连续的占用网格。 在这里,一个隐函数学习多个模板,这些模板可以经历另一个隐函数学习的非刚性变形。 与多图集分割 [21, 58] 一样,模板使我们的方法能够更好地处理有限的训练数据,并且对标签噪声不太敏感。 得益于直通估计器(STE)[3],可以以端到端的方式学习多个模板,而计算成本可以忽略不计。 最后,为了进一步提高数据效率,我们引入了卷积隐函数[9, 45]来提取多尺度特征。

为了证明 ImplicitAtlas 的有效性,我们对肝脏、海马体和胰腺的三个医学形状数据集进行了广泛的实验。 我们的方法明显优于当前基于隐式表示的方法 [12,38,45,72],特别是在几次学习设置中进行训练时。 我们还演示了几种潜在的应用,例如根据用户提供的点注释进行形状补全,以及通过学习的密集对应进行关键点标记。 正如将要展示的,尽管生物医学数据集带来了挑战,但所提出的隐式方法通过允许它们从训练数据中学习多个模板并选择一个来重建特定器官而变得非常有效。

2.相关工作

基于 CNN 的体积方法(例如 U-Net 及其变体 [10,22,40,53])现在在生物医学图像分割中占据主导地位。 这从 CHAOS 挑战 [29] 和 Medical Segmentation Decathlon [24] 的结果中可以明显看出。 两项比赛的获胜者都使用了依赖体积 CNN 的方法组合来处理传统的体积图像分割问题。 然而,生物医学图像分割在实践中还远未得到解决。 例如,保存几何和拓扑结构仍然很困难。 同时,可以准确地恢复体积:获得高质量的表面仍然很困难[34]。

这些挑战凸显了形状建模在医学图像分析中的重要性[16, 36],包括医学图像分割[33,50,55,59]、计算机辅助诊断[4,59,71]和计算解剖学[5] 。

显示表示。 由于医学成像中的形状通常在图像堆栈中的离散网格上进行注释(通过计算机断层扫描(CT)和磁共振成像(MRI)等获取),因此大多数现有技术[33,50,55,59,64]依赖于 体素表示。 作为 2D 像素网格的 3D 扩展,体素神经网络可以通过使用 3D 卷积扩展相应的 2D 版本(例如 3D UNet [11])或使用复杂的算子来开发 [67-69]。 不幸的是,它们对相对较低的空间分辨率具有较高的内存要求。 因此,还探索了几何数据结构(例如点云和网格)的使用。 点云在传感和处理方面轻量且灵活[7,23,48,70]。 它们适合提取语义信息[18,19,66],但不捕获拓扑。 从点云恢复表面是一个不小的问题。

三角网格允许对高保真表面重建进行内存高效处理[14,28,61,62],但改变其拓扑结构并非易事。 有一些为此目的而设计的算法[37],但它们需要特殊的启发式方法,而这些启发式方法不能很好地概括。

隐式表示。 最近,隐式表示 [8,38,44] 在基于深度学习的 3D 计算机视觉中变得越来越流行。 它们将 3D 形状表示为连续 3D 场中的等值面,并由深度网络参数化。 由于其灵活性、内存效率以及以任何分辨率表示任何拓扑的能力,隐式表示不仅在形状方面得到了广泛研究,而且在外观[39]和场景[42]建模方面也得到了广泛研究。

然而,他们尚未在医学图像分析方面取得重大进展。 该领域为数不多的研究之一可以在[49]中找到。 但它的重点是改进使用现有隐式表示方法产生的医学图像分割。 相反,我们专注于高质量的隐式表示方法,以解决生物医学成像领域开发隐式场的困难。

图谱和模板。 概率图谱广泛用于基于图集的图像分割 [21, 58],因为它们是处理生物医学图像中噪声的极好方法。 随着深度学习的出现,图谱已被集成到卷积神经网络中[2,13,20,54]。 所有这些方法都依赖于预先计算的图谱。 它们是通过融合多个手动注释的图像而创建的;图谱还必须预先注册到目标图像,以将其与感兴趣的结构对齐。 在[60]中,尝试使用可以自动对齐和变形以匹配目标结构的图谱。

另一方面,模板与[12, 72]中的隐式曲面结合使用。 这涉及使用隐式方法来预测模板周围的变形,其中变形和模板都是隐式参数化的。 然而,这些方法是在大型训练数据集上开发的,数据效率并不是主要关注点。 两者都依赖于 MLP 解码器,不会像卷积解码器那样引入空间还原偏差 [45]。 此外,这些方法只能自动学习单个隐式模板,而本文的一个中心论点是它对多个隐式模板都有好处。

3.Methodology

在本节中,我们首先简要回顾一下深层隐式形状表示。 然后我们介绍我们的模型、相应的架构和我们的训练方法。

3.1. 背景:深层隐式曲面

隐式形状表示 [8,38,44] 通过将 3D 坐标映射到形状指示器(通常是占用率或有符号/无符号距离)来建模形状。 在这项工作中,我们使用后者。 对于形状 S,该映射表示为

其中 h 是编码 S 的 c 维潜在向量,p(x, y, z) ∈ R3 是查询点,F 由输出占用概率 o ∈ [0, 1] 的深度网络实现。 对于 S 内的 p,o 应接近 1,否则应接近 0。 给定一个训练集,F和每个形状对应的向量h可以通过多种方式学习。 在这里,我们使用[38]中的自动解码方法,其中 h 被视为可学习参数并与 F 的参数联合优化。

除了直接输出形状指标的方法外,还有将形状视为模板变形的研究[12,72]。 F 重写为 T 和 D 的复合函数,即

其中 D : Rc × R3 → R3 是将查询点 p 映射到新坐标 p′ 的函数,T : R3 → R 是学习的隐式函数,其作用与等式 1 中的 F 相同。 1 但学习单一形状。 D 可以通过多种方式实现,包括加性变形 [12] 或逐点仿射变换 [72]。 在此公式中,T 扮演模板的角色,因为它对先前从训练形状中学到的形状进行编码。 值得注意的是,由于所有变形坐标都与模板对齐,因此可以轻松建立形状之间的密集对应关系。

3.2. 模型管道

尽管早于深度学习时代,多图集技术[21, 58]擅长处理有限数据和标签噪声。 我们将它们转化为我们的框架,如图 3 (a) 所示。 鉴于方程的表述。 2、我们取T为隐式模板网络(T),D为隐式变形网络,式(1) 2 被重写为

其中 t 和 d 是独立的向量,它们是 h 的函数,D 的输出是变形(而不是等式 2 中的变形坐标)。 请注意,T 现在有一个附加参数,我们将使用它从几个可能的模板中选择一个模板。 与早期的工作相比,这种设计能够学习多个模板,并提高了我们模型的表示能力 [12, 72]。

模板选择。 我们引入了一个可学习参数矩阵 T ∈ Rm×c,其中 m 表示模板的数量,c 是潜在向量的维度。 模板选择是通过从 T 中选取行向量 t 并馈送到隐函数 (T) 来实现的。 t 被取为

其中STE-Softmax是一种带有直通估计器的softmax[3]:softmax在前向传递中被“硬化”为one-hot,但它在后向传递中直接采用one-hot的梯度。 这可以被视为分类分布的重新参数化[30]。 STE-Softmax 可以被 Gumbel-Softmax [26, 35] 取代,理论上,它应该为分类重新参数化提供更平滑的梯度。 然而,在我们的实验中,这并没有什么区别。

估计变形。 方程的变量 d。 3 控制查询点 p 处的变形,我们将其视为逐点加性变形:p → p + D(d(h), p)。 由于医学形状以离散网格表示,因此 p 被实现为网格 P ∈ RD×H×W×3,其中 D ×H ×W 表示空间大小。

3.3. 网络架构

多层感知器(MLP)是参数化深度隐式函数的流行选择[8,12,38,44,72]。 不幸的是,这些 MLP 往往需要大量数据。 因此,我们使用了卷积解码器,如 ConvONet [45]。

如图 3 (b) 所示,潜在向量(在我们的方法中为 t 或 d)首先通过卷积层和上采样层转换为多尺度特征图,查询点 p 从中获取其特征作为以下函数: 它的坐标通过三线性插值[25]。 我们不像[45]中那样只使用最终的特征图,而是像[9]中那样从多个分辨率的特征图中提取多尺度特征。 最后,坐标以及编码局部和全局语义信息的结果特征被连接起来并输入到一个小型 MLP 中以产生输出。 多尺度特征使模型比纯 MLP 更少需要数据。

T 和 D 使用相同的卷积解码器实现,除了最后一层。 T 有一个输出通道,D 有三个输出通道。每个卷积块都是卷积层、组归一化 [63] 和泄漏 ReLU 激活 [65] 的堆栈。 第一个上采样层被实现为像素洗牌[52]。 更多实施细节可以在补充材料中找到。

3.4. Model Training

当我们采用占用公式时,要最小化的主要任务损失是二元交叉熵,

其中 N 是采样点数,o 和 ˆo 分别是预测占用率和真实占用率。 鉴于我们使用自动解码方法[44]来学习潜在向量并使用来自高斯先验分布的这些潜在代码,我们添加正则化惩罚

哪里 || · ||2 表示 l2 范数。 为了在使用有限数据进行训练时进一步规范化模型,我们定义了 2 个额外的正则化项:拉普拉斯平滑度 (LLS) 和变形惩罚 (LDP),写为:

其中,Fx 是 F(h, (x, y, z)) 的缩写,与 y 和 z 类似,D 是 D(h, p) 的缩写。 最小化 LLS 有利于空间平滑的输出,而最小化 LDP 会限制变形,以便 T 必须了解更多细节。

在训练过程中,我们统一采样分辨率低于感兴趣体积的网格 P∼ ε R ∼D× ∼H× ∼W×3。 对于本研究中 1283 的形状,我们在训练期间对 323 个网格进行采样。 网格中的每个点都添加了随机噪声,并且其占用地面实况在感兴趣的全分辨率体积中进行采样。 这种子网格技术显着降低了内存和时间方面的训练成本。 此外,当我们在统一网格上对查询点进行采样时,可以使用具有自定义内核的 3D 卷积来高效地计算 LLS。 LDP 是在网格上计算的。

4.Experiments

4.1. Datasets

在我们的实验中,我们使用医学分割十项全能 (MSD) [1],它可以说是迄今为止最大、最全面的医学图像分割数据集。 它是 10 个医学图像数据集的集合,其中包含使用多种不同模式成像的多个器官。 为了在许多不同的条件下测试我们的方法,我们对三个提出不同挑战的器官进行了实验:肝脏,一个具有复杂细节的大器官; 海马体,一个小器官; 胰腺是一种柔软的器官,可以以多种姿势看到。 由于注释仅适用于官方训练分割,因此我们将每个器官的数据分割为已知(训练)和未知(测试)形状。 所有器官注释均使用器官分割的边界框进行裁剪,并通过样条插值调整为固定大小 1283。 标签。 图1总结了3个结果数据集的基本信息。

4.2. 基线方法

当前大多数形状表示方法都是在 CAD 模型或场景上进行训练,而不是在生物医学形状上进行训练。 为了公平比较,我们重新实现它们并在我们的数据集上训练它们。 选择的网络大小与我们使用的网络大小相似。 我们在所有情况下都使用相同的训练和推理程序。 实施了以下基线方法

MLP 解码器。 由于 MLP 的简单性和灵活性,许多算法 [8,38,44] 使用它们作为解码器。 我们的 MLP 基线使用 DeepSDF 的架构 [44]。

MLP 解码器 + 模板。 最近依赖模板的算法 [12, 72] 使用不同的基于 MLP 的网络架构和变形公式。 为了使比较有意义,我们基于等式实施了基线。 3,其中 T 和 D 使用相同的 MLP 架构实现,这重复了[12]中所做的事情。

Conv Decoder.。 问题与依赖卷积隐式场的早期方法相同 [9, 45]。 为了进行比较,我们基于等式实施了基线。 1 使用与我们的解码器相同的架构。

4.3. 表示已知和未知的形状

我们首先介绍我们的实验过程,然后分析重建结果和消融研究。

实验设置。 与早期的工作[12,44,72]一样,我们首先评估模型对已知和未知形状进行编码的表示能力。 我们为肝脏、海马体和胰腺训练了三个独立的模型。 然后我们分别评估训练集和测试集的重建质量。 第一个被视为已知 (K) 形状的集合,第二个被视为未知 (U) 形状的集合。 对于前者,重建是在自动解码的训练期间进行的。 对于后者,我们优化随机初始化的潜在代码 h 以重建具有固定模型权重的形状。

除非另有说明,否则我们在方法中使用 m=5 模板。 一组固定损失权重用于所有情况:L2 为 10−3,LLS 为 10−3,LDP 为 10−2。我们在补充材料中提供了有关训练和重建过程的更多详细信息。

指标。 为了评估地面实况和预测重建之间的形状相似性,我们使用骰子相似系数(DSC)和归一化表面骰子(NSD)[43],它们是医学成像文献中使用的标准度量。 它们被定义为

消融研究。 我们进行了一项消融研究,以分析 ImplicitAtlas 各个组件的有效性。 我们在选项卡中报告结果。 3. 多模板(MT)提高了已知和未知形状的表示能力,特别是对于更困难的情况,例如胰腺。 正如之前所观察到的,方程的正则化项 LLS 和 LDP 如下: 7 增强对未知形状的泛化性能。 请注意,在某些情况下,单模板版本的性能低于没有模板的纯卷积解码器。 例如,比较选项卡中的“D”。 3 至选项卡中的“Conv Decoder”。 2. 这表明模板的性能改进取决于网络架构和数据集,如[12]中先前所述。

通过可视化,我们现在提供了 ImplicitAtlas + reg 的定性分析。 如上所定义。 形状重建和插值。 在图 4 中,我们显示了每个数据集 2 个随机选择的重建。 所有重建都是高质量的,即使表的数值指标也是如此。 尽管图 2 描绘了训练数据伪像,但图 2 并不完美。由于我们的隐式形状模型对训练样本的先验形状进行编码,因此重建比手动注释平滑得多,这表明我们的隐式形状模型的潜在用途 后处理人工注释的方法。 我们还显示了通过线性插值潜在代码获得的 2 个样本之间的插值。 插值样本看起来始终有效。 此外,该模型甚至捕获了姿势,如图底行所示。 这意味着我们的模型通过有限的数据学习丰富的语义。

学习模板和形状生成。 由于 T 的每一行都可以解释为潜在模板向量 t,因此我们可以通过使用 T(t,·) 创建占用字段来使用 T 来可视化它们。 我们对图 5 顶行使用的所有五个模板都执行此操作。在投影中,模板看起来相似,肝脏、海马和胰腺的平均成对 DSC (%) 分别为 96.10、94.89 和 94.15。 尽管如此,如表所示。 如图3所示,这些多个模板的存在以可忽略的计算成本提高了已知和未知形状的表示能力。 为了生成新的形状,我们从高斯分布中随机采样 h,计算方程 1 的 t(h) 和 d(h)。 3、将它们解码为占用字段。 可以通过调节前向传播中的 t 来实现从特定模板的生成。 在大多数情况下,生成的形状是多种多样且有效的,如图 5 的底行所示。但是,仍然存在诸如边界效应和孔洞之类的伪影。 更多的训练样本和复杂的数据增强技术有望解决这个问题。

4.5. Few-Shot Learning

实验设置。 由于医学形状通常稀缺,我们考虑数据集的极端设置:我们能否仅从 5 个训练样本中学习有效的形状模型? 为了测试这一点,在训练期间,我们仅使用每个器官训练集的前 5 个样本,同时保留第 2 节中使用的所有其他设置。 4.3.

结果。 我们在选项卡中报告结果。 4.我们的方法仍然明显优于基线,特别是在更具挑战性的胰腺病例中。 只有 5 个训练样本,MLP 的数据匮乏性质就变得很明显:在某些情况下它无法提供有意义的重建。 模板和卷积解码器都很有用,并且所提出的正则化项可靠地提高了重建精度。

5.应用

我们的模型引入的深度先验可用于生物医学成像中的众多下游应用。 在本节中,我们将探讨一些有希望的方向。

5.1. 从点注释完成形状

单击是一种为许多生物医学应用提供注释的广泛方法,例如交互式分割[51]、本地化[57]。 在这里,我们展示了 ImplicitAtlas 可以通过此类点击生成可接受的形状,从而使注释过程变得不那么麻烦。 在我们的实验中,我们在未知形状的器官边界附近采样了 64、256 和 1,024 个点。 随机初始化的潜在代码经过优化,可以通过最小化点的 LTask 损失来重建形状。 我们使用基线或 ImplicitAtlas 重复这些实验,并将结果绘制在图 6 中。通过将重建形状与 DSC 指标中报告的地面实况进行比较来评估形状完成性能。 我们的方法始终优于其他方法,尤其是在使用较少的点时。 我们在补充材料中提供了有关这些实验的更多详细信息。

5.2. Dense Correspondences

关键点和地标标记是医学成像中的一项重要任务[17,47,73]。 由于我们的隐式变形网络 D 的所有变形点都与模板很好地对齐,因此可以轻松建立多个形状之间的密集对应关系。 在图 7 中,我们突出显示了三个器官上的匹配关键点。 在每种情况下,我们在模板上手动选择 5 个关键点,并检查这些点的发送位置。 由此产生的对应关系在视觉上是可以接受的。 密集的对应关系可以应用于许多潜在的应用中。 例如,我们可以使用隐式函数创建“图集”:将模板标记为几个子部分,并将它们转移到所有形状注释中。 它还提供了一种模拟时空变化的方法,例如肿瘤生长 [32, 46]。

6.Conclusion and Future Work

我们已经证明,尽管在生物医学应用中存在挑战,但使用所提出的方法可以使深层隐式曲面变得有效。 未来的工作还有很多延伸方向。 首先,我们方法的当前版本适用于单个器官,而许多生物医学应用需要多个器官/部位。 将 ImplicitAtlas 扩展到医学成像领域的多类模型将特别有用,在医学成像领域处理不同姿势和尺度的多个对象将是一项挑战。 此外,由于隐式函数也可以用于外观建模,因此我们可以扩展ImplicitAtlas来对形状和外观的联合分布进行建模。 形状和外观的隐式模型将实现许多新颖的应用。 最后但并非最不重要的一点是,探索 ImplicitAtlas 的新应用将会很有趣。 模型编码的深度先验可以直接用于改进人类注释或模型预测。 它还可以作为医学图像分割、形状分析和多站点概括的工具。