SAMTooth:利用 SAM 实现稀疏监督下的牙齿点云细分 !

2024-09-11 19:39:16 浏览数 (1)

牙齿点云细分是许多正畸应用中的基本任务。目前的研究主要集中在全监督学习上,这需要昂贵且耗时的手动点逐一标注。尽管最近提出了使用弱标记进行3D细分和取得令人满意结果的弱监督替代方法,但在标记极其稀疏时,它们往往会失败。 受到Segment Anything Model(SAM)强大的提示分割能力的启发,作者提出了一种框架SAMTooth,利用这种能力来补充极度稀疏的监督。 为SAM自动生成适当的点促剂,作者提出了一种新颖的置信感知促剂生成策略,其中将粗分类预测与置信感知过滤聚合。此外,为了充分利用SAM输出中的结构和形状线索来帮助3D特征学习,作者提出了一个蒙版引导的表示学习方法,将SAM生成的牙齿口罩重新投影到3D空间,并约束这些不同牙齿的点生成具有区别的表达。 为了证明该框架的有效性,作者在公共数据集上进行实验,令人惊讶地发现,在只有0.1%的标注(每颗牙一个点)的情况下,作者的方法可以明显超过最近的弱监督方法,性能甚至可以与最近的完全监督方法相媲美,展示了将SAM应用于具有稀疏标签的3D感知任务的重大潜力。 代码在https://github.com/CUHK-AIM-Group/SAMTooth。

1 Introduction

准确地从口腔扫描仪(IOS)的网格数据中提取3D牙点的牙齿分割在许多正畸应用中起着关键作用,包括详细分析牙齿形态学、治疗计划、个性化装置设计等。然而,现有的牙齿点云分割模型 主要依赖于大量的标注数据集进行训练,这涉及到牙点云标注的劳动密集型任务。例如,一个有经验的牙科医生手动标注半个下颌大约需要15到30分钟 [30]。这种耗时过程对建立大规模、高质量标注数据集极具挑战性并且降低了诊断系统的泛化性 。

为了解决这个问题,人们对研究弱监督的替代方法越来越感兴趣。在不同的弱标签(涂鸦、框、部分点等)中,部分点是一个有前景的方向,因为标注的效率高,只需为每个牙齿标注一个或几个点。现有的基于部分点的技术从有限的标签中挖掘出各种训练约束,如扰动一致性 [33; 45],监控传递 [40; 7],自监督预训练,假标签等,这些方法已经显著减少了标注劳动力。然而,如图1所示,当标签稀疏度增加到0.1%(每个牙齿一个点)时,现有工作的最好成绩仅给出了 Baseline 6.18%的性能提升,与全监督正则化相比,产生了22.44%的mIoU差异,表明当标签极度稀疏时,现有工作无法表现良好。

作为处理此事的第一尝试,作者打算利用最近的SAM[10]模型。在110万张图像的数据集中训练,SAM可以给出由手动定义的视觉提示生成的细粒度 Mask 。如图1a所示,如果作者渲染输入3D模型的图像,并使用适当的提示将它们输入到SAM,作者可以得到每个牙齿的2D目标 Mask 。由于这些 Mask 包含明确的几何形状信息,作者可以利用它们来补充极稀疏监督。然而,直接将2D SAM应用于3D任务并非易事,因为存在两个问题。首先,很难自动提示2D SAM生成所需 Mask 。SAM Mask 的质量严重依赖于人类提供的提示,而模型训练期间整合人类输入是不现实的。其次,2D图像和3D点云之间存在显著差异,因此很难有效地利用SAM生成的2D Mask 来增强3D域中的模型学习 [26, 14, 25, 24, 12]。

为了解决这两个问题,作者提出了一种名为SAMTooth的新颖框架,用于极稀疏标签的正畸牙齿点云分割。如图1b所示,该框架包括两种模式,包括_自信提示生成_ (CPG) 和_ Mask 引导表示学习_ (MRL)。为了自动生成SAM可用于的适当提示,作者提出CPG将每个预测 tooth 的点聚并投影到图像平面。点预测可能会很噪声,所以作者对点进行信心度估计以过滤不可靠的聚合并选。为了充分利用 SAM 输出进行3D特征学习,作者进一步改进了 MRL 将 SAM 输出的像素重新投影回 3D 空间,并利用对比学习提供训练约束。考虑到背景点也需要受到约束,作者还计算了 SAM 目标 Mask 产生的背景 Mask 并施加显式监督。实验结果表明,SAMTooth 可以大大超过其他弱监督方法,甚至与使用0.1%标注的最近全监督方法相竞争。

2 Method

作者的框架是为了实现弱监督的牙齿点云分割,利用视觉基础模型SAM的零样本能力。如图2所示,它首先从输入扫描图像渲染和映射开始(第2.1节),将输入扫描图像渲染并建立3D点与2D像素之间的映射关系。然后,输入点云P传递到3D分割网络,得到粗略预测Y和点级置信度C,进一步传递到第2.2节的置信度感知提示生成,为SAM生成合适的点级提示。然后,SAM处理生成的提示和渲染图像,得到物体 Mask M,用于通过 Mask 引导的表示学习约束3D特征的第2.3节。整个框架通过分割约束和SAM输出的互补约束进行优化(第2.4节)。

Image Rendering and Mapping

为了利用SAM的输出进行3D表示学习,作者首先将3D IOS网格中的图像作为SAM的输入进行渲染。作者选择从网格而不是点云渲染,因为网格包含更多的纹理细节,而且在正畸应用中总是可用的。根据凸面相机的成像原理,每个点的投影坐标可以通过以下公式获得:

Confidence-aware Prompt Generation

SAM依赖合适的提示来生成高质量的物体口罩,这一过程将进一步影响后续的3D表示学习。因此,应谨慎设计一个合适的提示生成策略。在这方面,作者提出了一种自动提示生成策略,通过聚合每个粗糙预测的牙齿的3D点来生成提示,并伴有自信感知的筛选步骤,以消除那些会导致歧义的牙齿预测。

有了适当的提示点, SAM可以生成精确的目标 Mask ,从中作者可以发掘更多的约束以补充3D模型的稀疏监督。为此,作者提出将2D目标 Mask 重新投影到3D空间,并利用对比学习进行前景特征区分。考虑到背景点,即牙龈也应该受约束,作者进一步从前景 Mask 中计算出背景 Mask 并对其进行相应特征的正规化。

Model Optimization

3 Experiments

Experiment settings

数据集和评价

为了评估作者所提出的方法的有效性,作者在公共3DTeethSeg[1]数据集[1]上进行实验。牙齿识别遵循FDI世界牙科联盟的表示方法[1]。3DTeethSeg[1]是一个公开的牙齿分割数据集,包含900名患者生成的1,800个可用3D IOS扫描,根据实际患者年龄分布获取[1]。为了进行公平的比较,作者在所有实验中使用相同的划分方法:1,080个扫描用于训练,360个用于验证,剩余的用于测试[8]。遵循先前的牙齿分割方法[27, 28, 4],作者使用Jaccard指数(也称为mIoU),Dice相似系数(DSC)和点分类准确率(Acc)进行评估[1]。

实现细节

作者采用[43]中的标准ViT-B/16作为分割backbone。作者的框架使用AdamW优化器,学习率设置为5e-4,批处理大小为8,权重衰减为0.05[8]。作者实验中设置置信阈值为0.6,温度为0.1,预热周期为10,并将损失权重分别设置为8中的1/0.1/0.01[1]。参照先前的研究[28, 27],作者从IOS扫描中采样16,000个点构建输入点云,并使用三近邻插值策略在评估时将预测值上采样到原尺寸[1,27]。

Main results

为了与最近的最先进的研究进行公平的比较,作者使用相同的 Backbone 网络,并基于官方仓库重制他们的方法。作者将比较结果展示在1处。SAMTooth在mIoU上达到了76.47%,在mAcc上达到了86.64%,较之前的方法有显著提升。特别是,SAMTooth在mIoU上分别比II-Model[11],MT[33],Xu和Lee[40],PSD[45],和SQN[7]提升了15.32%,12.47%,6.47%,11.35%,11.82%,9.14%,和10.98%。当只有0.1%的标注时,SAMTooth可以实现与全监督 Baseline 相当的表现,这揭示了所提出框架的有效性,也显示了SAM在为有限标签的牙齿点云分割提供训练信号方面具有巨大的潜力。作者在图3中还提供了定性的比较。观察到作者的方法在边界区域(黑色 Box )的分割结果比使用其他弱监督方法进行训练的方法更好。

More analysis

3.3.1 Confidence-aware Prompt Generation.

为了评估CPG的有效性,作者尝试了另一种提示生成策略AGG,其通过每个子组的简单聚合获得点提示。如表2所示,这种简单的聚合策略会导致mIoU性能下降4.63%,表明需要信心引导。作者还报告了定性结果在图4(a)中,从中作者观察到AGG生成的提示倾向于从牙齿中心偏移,且这些提示可能导致错误的目标 Mask 。相比之下,CPG生成的提示通常位于牙齿中心附近,生成的 Mask 可以无缝覆盖每个牙齿,从而有利于后续的表示学习。

Mask 引导的表示学习除了MRL,作者还尝试了其他约束策略,包括仅使用前景和背景学习的FG和BG方法。如图4(b)所示,使用FG可以已经比 Baseline 获得12.40% mIoU的性能提升,因为前景特征学习的补充约束。同时,使用BG也可以带来3.68% mIoU的提高。此外,将FG和BG组合,即MRL,可以提高性能的最大15.32% mIoU的改善,揭示了MRL的有效性。

总结而言, Mask 引导的表示学习方法结合了多种约束策略,可以显著提高模型的性能,使得模型能够生成更精确、更符合预期的图像效果。

4 Conclusion

在这篇论文中,作者提出了一种新颖的框架SAMTooth,用于弱监督牙齿点云分割。该框架利用了最近的高级可提示基础模型,即SAM,来补充极为稀疏的监督(每个牙齿一个点)。

它采用了一种置信度意识化的提示生成(CPG),用于自动为SAM生成精确的提示,这些提示受估计的点级置信度的引导。

然后,它利用了 Mask 导向的表示学习(MRL),以充分利用SAM生成的细粒度 Mask 。

在两个基准上的大量实验表明,所提出的方法在很多方面都明显优于现有方法,展示了SAM在3D感知任务上的潜力。

与现有方法相比,所提出的方法在很多方面都明显优于现有方法,展示了SAM在3D感知任务上的潜力。

参考

[1].When 3D Partial Points Meets SAM: Tooth Point Cloud Segmentation with Sparse Labels.

0 人点赞