【技术追踪】TeethDreamer：从 5 张口腔照片实现三维牙齿重建（MICCAI-2024）

三维重建搞起来~

TeethDreamer：一种3D牙齿重建新框架，旨在恢复上下牙齿的形状和位置，引入大型扩散模型的先验知识和3D感知特征注意力机制，重建性能表现SOTA！

论文：TeethDreamer: 3D Teeth Reconstruction from Five Intra-oral Photographs
代码：https://github.com/ShanghaiTech-IMPACT/TeethDreamer（即将开源）

0、摘要

正畸治疗通常需要定期进行面对面的检查，以监测病人的牙齿状况。当面对面诊断不可行时，另一种方法是利用5张口腔内照片进行远程牙齿监测。然而，它缺乏三维信息，如何从这种稀疏视图的照片中重建三维牙科模型是一个具有挑战性的问题。

本文提出了一个名为 TeethDreamer 的三维牙齿重建框架，旨在恢复上颌牙齿和下颌牙齿的形状和位置。

给定5张口腔内照片，该方法首先利用一个大扩散模型的先验知识，生成具有已知姿态的新的多视图图像，以解决稀疏输入，然后通过神经表面重建建立高质量的三维牙齿模型。
为了确保生成视图之间的 3D 一致性，本文在反向扩散过程中集成了 3D 感知特征注意机制；此外，在牙齿重建过程中加入了几何感知法向损失，以提高几何精度。

大量的实验证明，该方法优于目前最先进的技术，具有远程监控正畸治疗的潜力。

1、引言

1.1、临床背景：正畸治疗

（1）正畸治疗重点纠正牙齿错位，需要长时间持续检测；
（2）口腔内扫描（intra-oral scanning）提供了一种获得高质量 3D 牙齿模型的方法，但通常既耗时又昂贵；
（3）使用智能手机拍照，方便又快捷，那么如何从 2D 照片重建 3D 牙齿模型呢；

1.2、传统方法的不足

（1）使用参数化牙齿模型重建：没有捕捉到牙齿的个性化细节，且预定义牙齿模板的质量也对重建结果有很大的影响；
（2）多视图立体声法（Multi-View Stereo, MVS）：需要校准的相机姿态来进行精确的重建，由于口腔照片的稀疏性和重叠小，目前的 structure-from-motion (SfM) 方法很难恢复准确的相机姿态；
（3）使用预训练扩散模型对自然物体进行三维重建：通常以单视图输入为条件来生成图像，这导致了看不见区域的重建不准确，提供的几何信息较差；

1.3、本文贡献

（1）采用预先训练的扩散模型，分割口腔内照片的牙齿图像，生成特定视点的彩色图像和相应的法线图；
（2）为确保不同视图之间的一致性，进一步从噪声彩色图像和法线图中构建 3D 感知特征，并在去噪过程中通过注意机制将其引入扩散模型；
（3）利用生成的彩色图像和法线图进行神经表面重建，重建三维牙齿模型。在重建过程中引入几何感知法向损失，提高了几何精度；

3D重建模型流程图：
在这里插入图片描述

2、方法

2.1、多视图跨域扩散模型

首先利用预先训练的 SAM 模型来分割前景牙齿区域；
如图2所示，使用包含上齿的四个分割图像作为模型输入，记为 ${x^{1:4}}$ ，且 $\in \mathbb{R}^{3×H×W}}$ ；

由于输入图像的稀疏性，难以重建高质量的三维牙齿模型，因此，选择借助生成扩散模型来增强观察视角。除了 RGB 图像，法线图（normal map）是恢复 3D 模型的另一个重要信息。
将分割的牙齿图像 ${x^{1:4}}$ 作为 Zero123 中预训练的扩散模型 ${f}$ 的输入条件，生成 ${N}$ 个预定义视点 ${v^{1:N}}$ 的彩色图像 ${c^{1:N}}$ 和法线图 ${n^{1:N}}$ ，记为：
在这里插入图片描述
TeethDreamer总览：

${c}$ 和 ${n}$ 与 ${x}$ 具有相同的维度，真实彩色图像和法线图是由配对的口腔内扫描模型预合成的，如图1所示。这样可以利用扩散先验较强的 zero-shot 泛化能力，此外，还可以利用法线图中丰富的几何信息来提高牙齿重建的精度；

目的是学习所有这些视图的联合分布 ${p_θ(c^{1:N},n^{1:N}|x^{1:4})}$ ，它可以用数学形式表示成一个多视图扩散模型：
在这里插入图片描述
如图 2(a) 所示，将输入视图 ${x^{1:4}}$ 与带噪目标视图 ${\{c_t^i, n_t^i\}}$ ， ${i=1,...,N}$ 拼接在一起，作为 UNet 的输入，此外，与 Zero123 相同，本文还使用 stable diffusion 的注意层（图2(a)中的文本注意分支）处理目标视点 ${v^i}$ 与输入视图 ${x^{1:4}}$ 的 CLIP 图像特征的拼接。

由于输出通道数量的差异，训练扩散模型同时生成彩色图像和法线图会影响预训练模型的性能，为解决此问题，采用一个域切换器 $\in \mathbb{R}^{1}}$ ，它决定了输出类型是彩色图像还是法线图，训练了一个多视图跨域扩散模型，将公式 (1) 修改如下:
在这里插入图片描述

2.2、3D 感知特征注意

保持从不同视图生成的图像和法线图的一致性，对于高质量的几何重建至关重要，本文引入了一种采用 3D 感知特征提取和深度注意机制相结合的方法。

该策略集成了当前目标视图 ${\{c_t^i, n_t^i\}}$ 去噪过程中的中间状态 ${\{c_t^{1:N}, n_t^{1:N}\}}$ ，首先将潜空间中生成的 2D 图像 ${c_t^{1:N}}$ 和法线图 ${n_t^{1:N}}$ 反向投影到预定义的尺寸为 $64^3$ 的 3D 体素网格上，利用 3D CNN 分别对颜色和正常特征体进行编码，随后，3D U-Net 将这些特征体合并，确保输出在几何形状和外观上保持一致。

为了提取目标视点 ${v^i}$ 的特征，创建了一个视锥体，并在得到的 3D 特征体内执行插值。通过深度注意层将这些特定于视图的特征集成到去噪过程中。该方法有效地捕获了不同视图之间的空间关系，整合了目标视点的基本信息，显著提高了生成视图之间的一致性。（感觉没有描述的很详细~）

2.3、几何感知牙齿重建

首先从生成的彩色图像 ${c^{1:N}}$ 或法线图 ${n^{1:N}}$ 中分割牙齿掩模 ${m^{1:N}}$ ，然后从训练集 ${\{n^{1:N},c^{1:N}, m^{1:N}\}}$ 中随机抽取一批具有相关光线的训练像素 ${\mathcal{R}_k = \{n_k, c_k, m_k, d_k\}∈R}$ ，用于神经表面渲染。其中， ${n_k∈\mathbb{R}^3}$ ， ${c_k∈\mathbb{R}^3}$ ， ${m_k∈\{0,1\}}$ ， ${d_k∈\mathbb{R}^3}$ 分别表示第 ${k}$ 条光线的法线值、颜色值、掩码值和光线方向，整个目标函数定义如下:
在这里插入图片描述
其中 ${\mathcal{L}_{normal}}$ 表示发现损失项；
${\mathcal{L}_{rgb}}$ 测量渲染颜色 ${\hat{c}_k}$ 和生成颜色 ${c_k}$ 之间的差异；
${\mathcal{L}_{mask}}$ 计算渲染掩码 ${\hat{m}_k}$ 和分割的掩码 ${m_k}$ 之间的二进制交叉熵；
${\mathcal{R}_{eik}}$ 是为了保证重建表面的平滑性而进行的正交正则化，而 ${\mathcal{R}_{sparse}}$ 是为了减少浮点数而进行的稀疏正则化， ${λ}$ 是每个损失项对应的权重项；

下面深入探讨 ${\mathcal{L}_{normal}}$ ，利用 Neus 中 SDF 表示的可微特性，可以通过计算 SDF 的二阶梯度得到固有重构曲面的正态值，采用几何感知的法向损失来最小化所呈现的法向 ${\hat{n}_k}$ 与参考 ${{n}_k}$ 之间的差异：
在这里插入图片描述
${w_k}$ 是一个几何感知权值，定义为:

${ϵ}$ 是一个接近于零的负阈值。这种设计背后的基本原理是，视图方向 ${d_k}$ 总是与表面法向 ${{n}_k}$ 相反，即它们之间的夹角应该始终落在 ${[90^◦,180^◦]}$ 的范围内。任何偏离此条件的情况都表明生成的法向 ${{n}_k}$ 不准确，从而减少了对表面重建过程的指导作用，即 $w_k = 0$ 。

3、实验与结果

3.1、数据集

（1）3200例，每例5张图像，3000训练，100验证，100测试，调整为256×256大小；
（2）Blenderproc方法在八个特定视点合成口腔内扫描的目标彩色图像和法线图，合成图像256×256大小；
（3）对于仰角，下牙的视点一半设置在45度，另一半设置在67.5度，上牙则相反，方位角分为四组：30度，60度，120度和150度；

3.2、实施细节

（1）扩散模型微调阶段：3D CNN、3D UNet和注意层是可训练的，而扩散 UNet 和 CLIP 被冻结；30000步，batch size=64；学习率 $1e^{-5}$ ，前10000步，线性增加到 $5e^{-4}$ ；
（2）3D 牙齿重建阶段：Neus模型训练20000步，光线批处理为4096，学习率 $1e^{-5}$ ，前500步，线性增加到 $5e^{-4}$ ，然后指数衰减到 $5e^{-5}$ ；
（3）显卡：一个 A100 GPU；