APE vs RPR,改进3D图像检测,在标志点检索和定位任务中达到 SOTA性能 !

科技   2024-11-25 09:00   上海  

点击下方卡片,关注「集智书童」公众号

点击加入👉「集智书童」交流群




想要了解更多:

前沿AI视觉感知全栈知识👉「分类、检测、分割、关键点、车道线检测、3D视觉(分割、检测)、多模态、目标跟踪、NerF

行业技术方案👉AI安防、AI医疗、AI自动驾驶
AI模型部署落地实战👉「CUDA、TensorRT、NCNN、OpenVINO、MNN、ONNXRuntime以及地平线框架」

欢迎扫描上方二维码,加入「集智书童-知识星球」,日常分享论文、学习笔记、问题解决方案、部署方案以及全栈式答疑,期待交流!

免责声明
凡本公众号注明“来源:XXX(非集智书童)”的作品,均转载自其它媒体,版权归原作者所有,如有侵权请联系我们删除,谢谢。

作者提出了一种自监督模型,该模型可生成个体医学图像 Voxel (voxels)的3D解剖位置编码(APE)。

APE编码 Voxel 的解剖邻近性,即相同器官或邻近器官的 Voxel 总是比远离身体部位的 Voxel 具有更接近的位置编码。

与现有的解剖位置编码模型不同,作者的方法能够有效地为整个体积输入图像生成 Voxel 级嵌入映射,使其成为不同下游应用的优选方案。

作者在8400张公开可用的腹部和胸部CT图像上训练作者的APE模型。作者对该模型在解剖标志物检索和弱监督的13个腹部器官的少样本局部化等方面的现有模型比较。

作为一个实际应用示例,作者展示了如何以0.99的召回率廉价地将APE模型用于裁剪不同感兴趣解剖区的原始CT图像,同时将图像体积减少10-100倍。

代码和预训练的APE模型在https://github.com/mishgon/ape。

1 Introduction

人体具有3D几何结构。在大多数病人中,不同身体部位的相对位置大致相同。这激发了所谓身体部位回归模型,这些模型旨在学习医疗体积图像中不同部位的3D位置编码。学习这些模型的主要思想是基于其外观预测图像块/切片/ Voxel 在3D物理空间中的相对位置。以这种方式预训练的网络可以 later 作为少样本学习器进行器官感兴趣区域定位,或应用于图像配准、检索或病变跟踪。

无监督人体部位回归器(BPR)[11]和Deep-index [8]提出了CT图像中人体部位回归的算法,输出每个体积图像的2D横截面上一个单数字。BPR使用两个组成部分的损失:一个部分是两个轴层次之间的距离,另一个部分是符号。然后,Deep-index简化损失为两个轴层次之间的有符号距离,从而证明类似的位置性能。这些方法虽然快速,但提供粗糙的标注,只允许粗略的定位,如“肺位于横截面X和Y之间”,并对小器官不准确。

相对位置回归(RPR)[6]可以从1D到3D看作是对[11]和[8]的扩展。作者训练了两个ResNet类似的模型("粗略"和"精细")来预测贴片的3D位置嵌入,使得嵌入之间的偏移量与贴片中心在mm.s处之间的偏移量相匹配。他们方法的主要缺点是采用了一个贴图到向量的体系结构,因此无法构造一个 Voxel 级嵌入图。

自监督解剖学eMbedding(SAM)[9]是一种对比学习方法,用于自监督地学习 Voxel 级的表示,捕获全局和局部的语义。它输出两个高/低分辨率的特征图,包含包含多尺度高维 Voxel 级嵌入,作者将其用于解剖学标志物检索任务。产生地图的高维性是图像检索任务的一个缺点,因为它具有高的存储内存占用:一个SAM图需要比原始图多256倍的存储空间。

在作者的工作中,作者提出了一种新模型,能产生包含个体 Voxel 3D解剖位置嵌入的高分辨率图。与RPR[6]不同,作者的模型能够为所有个体 Voxel 产生位置嵌入,由于其具有UNet样式的架构。作者不是使用两种不同的粗糙和细粒度模型,而是在作者的新训练策略的启发下,通过对比方法[9],确保作者的位置嵌入捕获全局和局部解剖学语义。因此,作者的方法在[6]的块级多阶段方法与[9]产生高度 Voxel 级嵌入的直接端到端对比方法之间建立了桥梁,见表1。

作者的关键贡献有两个。首先,作者提出了一个能够产生包含个体 Voxel 3D解剖位置嵌入的高分辨率图的新模型,具有自然产生的属性:对 Voxel 的物理坐标的连续性;嵌入维度的正交性;对输入图像块的等效力,以及其宽度的变化。其次,APE在弱监督的少样本器官定位方面展现出优越性能,实现了总定位召回率,同时减少了图像体积的10到100倍。

2 Method

APE model

APE的核心思想是学习在"平均患者"的3D物理空间中,各种放射学特征之间的相对位置。作者通过从同一张断层扫描图像中采样不同的 Voxel ,并基于其局部特征将其嵌入到3D位置嵌入空间来实现这一想法。主要的训练目标是使预测的位移嵌入与对应 Voxel 之间的物理距离高度相关。尽管物理距离仅定义为同一 Voxel 中的 Voxel 之间的距离,但是,不同患者相似的解剖标志点的APE嵌入之间仍能够很好地对齐。因此,称它们为解剖位置嵌入。

以下,作者将详细介绍APE架构及其训练过程,如图1所示。

架构与RPR不同,APE模型可以有效地为输入图像的所有单个 Voxel 生成嵌入。这是通过类似于UNet的[4]架构实现的:它接受一个体积图像块大小为作为输入,并预测一个大小为的张量,作者称其为APE映射。该映射包含大小为的单个 Voxel 嵌入张量,称为APE嵌入。为了生成整个CT图像的APE映射,作者采用分块预测的方式。

关于网络结构,有两点重要改进:

1)为了使推理快速且内存高效,作者在第一层使用步长为的卷积,最后一层的APE映射通过一个因子为4的trl线性上采样;

2)在最后上采样层之前,作者使用没有缩放变换的批标准化层来确保APE嵌入在结构层面上的均值为0,标准差为1。

naive训练过程在基本形式中,一个APE训练迭代如下所示。作者从同一图像中样本一个批次的体积块,并从中采样个 Voxel ,总共有个 Voxel 。

然后,作者预测它们的APE嵌入,并计算对应的欧几里得距离. 作者还计算每个 Voxel 的绝对坐标相对于原始图像,并归一化到零均值和单位标准差,以便与预测的嵌入对齐。作者得到 Voxel 之间的目标对欧几里得距离. 训练目标是:

作者将这个最基本的APE版本称为_APE-naive_。请注意,如果 Voxel 属于不同的独立 Patch ,那么距离通常相对较大,所以作者称它们为_global_。相反,如果 Voxel 属于相同的 Patch ,距离相对较小,作者称它们为_local_。

由于网络将整个 Patch 作为输入,因此不需要依赖其外观来预测其 Voxel 之间 local 的距离。因此,作者认为_APE-naive_表示仅在全局层面上与解剖位置对应,而在局部层面上的解剖标志物对齐较差,作者在第3节中进行了实证验证。

改进的训练过程为了实现APE表示的预期local属性,例如在局部尺度上与解剖标志物对齐更好,对不同图像裁剪的等同性, Voxel 间距的变化以及颜色失真,作者提出了一种改进的训练过程。受到 Voxel level表示的对比学习启发[9, 3],作者采样对重叠增强 Patch ,并从中采样个正交 Voxel 对,即具有原始图像中相同绝对位置 的来自不同 Patch 的 Voxel 对。

Patch 增强包括缩放至随机间距,随机 Mask ,随机高斯模糊/锐化/噪声,以及截断强度至随机Hounsfield窗口。然后,作者预测所有个正交 Voxel 对的APE表示,并计算对间距离。目标距离与之前相同,其中是绝对正交 Voxel 对的归一化坐标

注意, Voxel 来自同一重叠区域之间的距离是local的,但是现在预测的距离是来自不同 Patch 的表示物之间的距离,这阻止了简单的解决方案,并在局部尺度上提高了APE属性。修改后的训练目标是

第二项 额外施加了对 crops 和 patch 扩充的 APE 嵌入的等同性约束。在作者的实验中,通过比较具有 的小写字母(APE-Aug)和具有 的大写字母(APE-equiv) 的模型,来消融其效果。

Evaluating APE on anatomical landmark retrieval

根据[9]的研究,作者评估APE嵌入如何对应于解剖位置。为此,作者在训练图像上标记一个与特定解剖标志相对应的 Voxel ,并计算其APE嵌入,作者称之为_query_。为了在测试图像上检索到与感兴趣解剖标志的近似解剖位置最近的 Voxel ,作者计算其APE图,并在该图上进行最近邻搜索,以查找与 Query (通过欧几里得距离)最近似APE嵌入的 Voxel 。

作为质量度量,作者计算径向误差(以mm为单位),即检索到的 Voxel 与测试图像中感兴趣解剖标志的 GT 位置所对应的标记 Voxel 的物理距离。作者报告了所有数据集中的平均径向误差(MRE)和径向误差的标准差。

Evaluating APE on few-shot organ localization

遵循 [6] 中的研究,作者在一个更具实践导向的应用中评估了 APE - 弱监督的少样本器官定位。在这个任务中,作者的目标是构建一个针对特定器官的边界框预测器,该预测器基于少量的训练图像(例如,边界框或关键点)进行廉价标记(如边界框或关键点)。

将 APE 应用到这个任务的想法是定义器官的一些解剖学 landmarks(来源器官边界框的构建),将边界框预测问题简化为这些 landmark 检索。为了与 [6] 进行公平比较,作者使用相同的框定义 landmarks 选择——器官的六个边缘点(沿图像轴线每个有两个点)。在训练阶段,作者在一张训练图像上对它们进行标记并预计算其 APE 嵌入。在推理阶段,作者使用预计算的每个嵌入作为 Query ,在测试图像上检索最接近解剖位置的 voxel,如第 2.2 节所述。然后,预测是基于检索到 voxel 的边界框获得的。为了实现少样本学习,作者对所有训练图像重复该过程并将测试图像上的边界框预测取平均值。

将 IoU(预测框和地面真相器官框之间的交集/并集)作为质量度量。此外,作者将评估当用来裁剪到感兴趣器官的原始图像作为降低进一步图像分析步骤的计算负担的预处理步骤。在这个场景中,达到几乎完全召回至关重要,以免裁剪掉器官的一部分。最简单的提高召回的方法是相对于中心的大小将预测框放大 倍。对于每个器官,作者将选择 值,使其相对于器官 Mask 的召回不低于 0.99,并计算原始图像的体积和裁剪到放大框的图像的体积之间的比率。作者将这种度量称为 0.99 召回时的体积比率(VR@99)。作者报告在测试图像上的 IoU 和 VR@99 的平均值和标准差统计数据,以及用训练轮次五份(5-shot 学习)进行交叉验证的结果。

3 Experiments & results

本研究旨在评估不同机器学习算法的性能差异。作者在多个数据集上进行了实验,并对结果进行了分析。实验结果表明,支持向量机(SVM)算法表现最好,而朴素贝叶斯分类器(Naive Bayes Classifier)和随机森林(Random Forest)算法的性能也相对较好。在对比实验中,作者使用了多种分类器和特征选择技术,并对它们进行了比较。本研究的结果可以帮助人们更好地了解机器学习算法的性能差异,并为选择合适的算法提供指导。

Datasets & implementation details

作者使用三个公开的CT数据集对APE模型进行训练。作者将来自AMOS的2400个CT和FLARE2022 未标注 部分的2000个CT合并到腹部域,同时还包括NLST的4000个胸部CT。所得的训练集完全为无标签数据。该训练集涵盖了腹部和胸腔域,大致比例相等,为模型在两个区域中的稳健性奠定了基础。

作者使用APE模型对每个批次的对重叠 Patch 进行训练,每对 Patch 包含对体积,训练将持续150k个批次。在单个NVIDIA RTX A4000-16GB GPU上进行训练,需要17小时。作者使用带有固定学习率0.0003、权重衰减和梯度裁剪为1.0范数的AdamW优化器。预处理仅包括裁剪以得到密集的前景体积(通过HU阈值进行阈值处理)。 Voxel 间距在 mm³的范围内随机增加。 Patch 大小增加,使 Patch 体积为个体积,且截面比例不超过2。

为了评估APE和 Baseline ,作者使用FLARE2022的标记部分。它包括50个CT扫描和13个腹部器官的分割 Mask 。

Anatomical landmark retrieval

在本研究中,作者对两种类型的解剖学标志物进行了评估:器官质量中心(mass centers)和器官边缘点(edge points)。作者分别将其对应的平均径向误差定义为,并在表2中报告了推理时间。注意,_APE-Aug_ 在 APE-Naive 和 APE-Equiv 之上,证明了作者在第2节中描述的直觉。与RPR [6] 相比,_APE-Equiv_ 通过使用更为简单的单阶段检索过程取得了相同性能,而RPR 依赖于蒙特卡洛集成,并使用具有独立粗粒度和细粒度模型的两阶段流水线。令人惊讶的是,SAM 在作者的实验中表现不佳,尽管作者在原始论文 [9] 中报告了非常好的检索结果。作者使用了他们的预训练权重和代码,因此作者的最佳猜测是性能下降是由于FLARE2022与 SAM 训练数据的域转移。

Few-shot organ localization

表格3、4中展示了作者的APE模型和 Baseline 的少样本器官定位IoU和VR@99。作者采用了nnUNet [4]和nnDetection [2]作为标准的监督 Baseline 。请注意,作者假设训练集中只有器官的弱标签。因此,在训练nnUNet时,作者使用填充边界框作为 Mask 。 Baseline nnUNet和nnDetection的表现都比APE差,作者解释为他们较弱的少样本能力。

能够生成1D解剖位置的位移嵌入的模型 [11, 8]结果较差,因为他们只能够预测沿单一轴的边界框。

作者的_APE-Equiv_模型在包括SAM [9]和RPR [6]在内的其他所有模型中表现突出,Wilcoxon符号秩和检验的p值小于10^-6。

定性结果

图2显示一张图像及其在轴向、冠状和矢状投影下的APE图的三个通道。尽管是分块推理,但APE嵌入与 Voxel 的物理坐标是连续的。

作者通过构建FLARE2022 [7] 中50名患者的13个腹部器官中心 Voxel 的APE嵌入的3D散点图,可视化不同器官和解剖标志在3D APE嵌入空间中的位置,如图3所示。请注意,这些嵌入结果是聚集的,并且簇的相对位置类似于腹部器官的相对位置。

4 Conclusion

在这项工作中,作者提出了一种3D解剖位置编码模型。实验结果显示,APE在解剖标志点检索和少样本器官定位任务中均达到了SOTA水平。

唯一与APE竞争的现有模型是RPR [6]。然而,与RPR不同,APE能够为输入图像中的所有 Voxel 生成嵌入,这简化了解剖标志点检索流程,使APE在其他应用场景中更加方便和高效。

当前APE模型存在几个限制。首先,作者尚未使其具有翻转和旋转的等效性。其次,它们仅在腹部和胸部图像上进行训练,并在腹部进行评估。

参考

[1].Anatomical Positional Embeddings.

扫码加入👉「集智书童」交流群

(备注:方向+学校/公司+昵称



点击下方“阅读原文”,
了解更多AI学习路上的「武功秘籍」


集智书童
书童带你领略视觉前沿之美,精选科研前沿、工业实用的知识供你我进步与学习!
 最新文章