论文标题
TANet：迈向全自动排牙
Wei G, Cui Z, Liu Y, et al. TANet: Towards Fully Automatic Tooth Arrangement[C]//European Conference on Computer Vision. Springer, Cham, 2020: 481-497.
摘要
确定最佳的目标牙齿排列是数字化正畸治疗计划的关键步骤。用于指定目标牙齿排列的现有实践涉及繁琐的人工操作，并且操作结果的质量严重依赖于个体专家的经验，导致治疗结果无效和出现未知的变化。本文提出了一种基于学习的快速自动排牙方法。为了实现这一点，我们将牙齿排列任务公式化为一个新颖结构化6-自由度姿态预测问题，并且通过提出一个新的神经网络结构来解决这个问题，该神经网络结构从大量的成功编码的正畸治疗案例的临床数据中学习。我们的方法已经通过大量的实验得到了验证，并且质量和数量上都显示出了令人满意的结果。
关键词
深度学习（Deep Learning），畸齿矫正学（orthodontics），牙齿排列（tooth arrangement），6D 姿态预测（pose prediction），结构（structure），图神经网络（graph neural network）
Sec01 介绍
不规则的牙齿排列不仅会导致美学问题，还会损害咀嚼的功能。不正确的咬合关系（如：过度咬合或者牙齿拥挤）可能会导致咀嚼障碍，这往往会诱发其他继发性疾病。随着人们对口腔健康的日益关注，对正畸治疗的需求也越来越多。尽管寻求正畸护理的人数正在迅速增加，但是能够满足需求的合格的正畸医生在总体上严重缺乏。目前，正畸治疗涉及繁琐的手工操作，并且培训专业的正畸医生也是一个漫长而且昂贵的过程。此外，诊断和治疗的质量在很大程度上取决于正畸医生的个人技能和经验。因此，有必要开发一个全自动的系统来快速推荐最佳的牙齿排列，以提高正畸治疗计划的质量与效率。
牙齿排列是正畸治疗中必不可少的一步。给定患者的一组位置不良的牙齿，牙齿排列的目标是预测理想的牙齿布局，这个布局作为通过正畸治疗实现的最终排列。为了生成令人满意的排列，需要考虑许多的因素。这使得牙齿排列成为一项复杂的任务，任务结果的质量严重依赖于正畸医生的专业技能和主观判断。
现有的在正畸治疗计划中使用的计算机辅助系统提供了用户界面，可以对单个牙齿手工编辑和可视化。作为口腔修复学的相关工作，Dai[^8]根据一套启发式规则进行了全口义齿的牙齿排列，牙齿从预先指定的牙齿组中选出。相比之下，我们打算解决一个不同的，并且更有挑战性的牙齿排列问题，问题与正畸治疗中病人的特定齿列的排列相关。在[^6]的工作中通过将上牙和下牙视为两个刚性物体来自动建立适当的牙齿咬合，而我们的牙齿排列问题需要调整齿列中每颗牙齿的姿势。
对于每个特定的患者，自动地确定牙齿的理想位置是极具挑战性的。尽管像“Andrew六个关键问题”[^1]这样的临床规则提出了牙齿正确排列的必要条件，但是患者牙齿的实际布局可能会对达到理论上的理想姿态产生妨碍。因此，通过基于规则的方法开发的数学模型很难在临床上导出可行的结果。除此之外，在牙齿模型上检测特征点或者其他人类定义的特征是一个乏味的过程，并且可能在开始姿态预测时引入误差。此外，牙齿模型的纹理较小，并且缺乏清晰的特征，尤其是当我们只有牙冠（牙龈外的牙齿）时。这些特征使得精确地和一致地定义牙齿的方向、位置或者其他低级特征变得困难，尽管这些特征是基于规则的方法的先决条件。
我们提出了一种基于学习的方法，从患者治疗前的初始的不规则的牙齿位置预测最佳治疗目标，从而开发了第一种用于正畸治疗的自动牙齿排列方法。我们将牙齿排列任务公式化为一个结构化的6D姿态预测问题，这个问题尚未被计算机视觉领域充分探索。我们的网络旨在通过监督学习来近似表示一种映射关系，这个关系是初始牙齿排列的牙齿输入模型到理想目标姿态的映射。这个网络由四个主要组件构成：用于颌骨级和牙齿信息的特征编码模块、用于牙齿之间信息传递的特征传播模块、用于6-自由度姿态预测的姿态回归模型和用于刚性变换的可微的牙齿装配模型。网络的损失函数是专门设计的，用于捕捉不同排列之间的内在差异，强化紧凑的空间关系，并且模拟基准数据中的不确定性。
总之，当前工作的主要贡献是：
- 我们开发了第一个基于深度学习的自动排牙框架；
- 我们提出了使用基于图的特征传播模块来更新由PointNet提取的特征，从而为成功地解决由牙齿排列任务引出的结构化姿态预测问题提供关键的上下文信息 。
- 我们提出了一种新的损失函数，函数基于排列错乱的牙齿布局的分布构成，它能够通过这个分布的内在差异、空间关系和不确定性，为牙齿对齐提供有效的监督。
Sec02 相关工作
6-自由度姿态估计问题
近些年来，姿态估计问题得到了广泛的研究。它旨在推断存在于RGB图像[^3],[^5],[^7],[^18],[^25],[^33],[^34],[^45]，RGB-D图像[^35],[^39],[^40]，或者点云数据[^26],[^29],[^30],[^44]中的物体的三维姿态，这个姿态具有六个自由度。现有的方法大致可以分为物体坐标回归法和模板匹配法。基于坐标回归的方法假设对应的3D模型对于训练是已知[^36]的，在这个假设条件下方法在像素级估计对应于每个物体的物体表面。基于模板匹配的方法使用各种技术（例如：迭代最近点[^4]，Iterative Closest Point, ICP）在已知的三维模型和图像观察之间执行对准。所有这些先前的工作都没有考虑多个对象及其相对关系，而我们面临的牙齿排列问题需要预测所有牙齿的6-自由度姿态（即多个对象）在相同时间形成规则布局。最重要的是，6-自由度姿态估计问题涉及到已知的3D形状的姿态与其图像观察之间的关系。相比之下，我们的问题是预测排列整齐后的牙齿的姿态，我们的目标通过学习正畸治疗的临床数据来解决这个更具挑战性的问题。
家具布置或者布局问题
最近有许多关于如何自动生成由各种家具对象组成的优化了的室内场景的研究[^10],[^14],[^21],[^37],[^38],[^42]。为了简化问题，大多数的这些方法使用边界框作为代理来粗略地近似输入对象，而不考虑对象的细粒度的几何细节。[^42]使用已知的先验去优化给定的三维模型的配置。他们方法的核心是用一组启发式规则来定义能量函数。[^31]中的方法解决的是一个3D对象相对于其他对象的放置问题，文中假设所有的对象都以相同的方向预对齐，因此只需要预测新添加组件的平移量。由于人造的家具形状通常具有明显的尖锐特征，因此可以很容易地确定这些物体的方向。作为比较，我们考虑的牙齿模型缺乏这种明显的特征，这使得精确地定义方向变得困难。此外，大多数家具布置问题的研究工作试图为给定的室内场景生成不同的布置，而牙齿排列问题的目标是每个特定患者的最佳牙齿排列。
3D形状生成问题
3D形状生成问题旨在根据用户规范或者通过从图像或者部分模型中的推断来生成逼真的三维形状。条件生成方法[^22]可以基于输入的条件生成逼真的图像。随着几何学习和三维表示方法的进步，各种生成模型被提出作为处理三维形状的强大工具[^24],[^27],[^28],[^32]。有条件的三维形状生成问题和自动排牙问题都旨在根据给定的条件生成三维形状。条件三维形状生成[^11],[^13],[^17],[^23]的最新工作集中于生成能够适应不同形状变化的逼真的结构化形状。然而，它们没有保留输入对象的几何形状，而这是牙齿排列问题中的一个硬约束。
Sec03 方法
3.1 综述
图1：我们方法的全部管道和网络架构。在第一阶段，输入的3D牙齿模型被自动分割，以产生每个牙齿的标签和点云表示。然后，对每颗牙齿的点集进行归一化和采样。第二阶段，网络组成的4个组件：特征编码、特征传播、转换（即：姿态）回归和牙齿装配模块。最终输出是一个重新的排列的齿列。为了清楚起见，这里仅绘制了牙齿级别的编码器。
如图1所示，我们提出的方法包括了两个主要阶段。首先是预处理阶段，从整个模型中分割出牙冠，然后对每个牙冠进行语义标注。第二阶段使用具有四个主要组件的网络来执行以下功能：
- 一组基于PointNet的点特征编码器，用于颌骨级和牙齿级的特征提取；
- 一个基于图的特征传播模块，其在牙齿之间传递信息；
- 每个牙齿的回归器用于组合其对应的牙齿级特征、全局特征和随机条件向量，并将组合结果作为输入，并且输出相对于该牙齿的输入位置的6D变换；
- 装配器，用于将表示在轴角表示中3D旋转映射到用于变换点的旋转矩阵中，并且输出重新排列的点云。
这些功能的细节将会在后续的小节中描述。
3.2 预处理
分割和标记作为牙齿排列算法的预处理操作至关重要。在三维牙齿网格上，精确的自动语义分割和标记存在许多现成的方法[^20],[^41],[^43]。我们使用[^41]中的方法，牙齿的标签是根据FDI两位数表示法来标记恒牙。请注意，我们只保留所有牙齿的牙冠用于牙齿排列中的计算。然后，用这些冠组成的模型定义一个局部坐标系，通过将其与世界坐标系粗略对齐来标准化位置和方向。得到的牙齿集合表示为$X={X_v\subseteq\mathcal{R}^3|v\in\mathcal{V}}$，其中$\mathcal{V}$是牙冠标签的集合，$X_v$是拥有标签$v$的牙冠的点云。
3.3 网络
牙齿居中
因为输入的牙齿稀疏地分布在空间中，这可能增加了彼此距离相距较远的牙齿之间捕捉特征的难度，所以我们首先把所有的牙齿平移到原点得$\tilde{X}_v={p’=p-c_v|p\in X_v}$和$\tilde{X}={\tilde{X}_v|v\in\mathcal{V}}$，$c_v\in C$是牙齿$v$的几何中心。这个措施是将牙齿的中心位置从其他特征中分离的关键，从而使编码器能够采用与平移无关的方式专注于提取几何特征，例如：形状的细节、方向和尺寸。
特征编码器
我们网络中的特征编码器是基于PointNet[^27]。使用对称函数，PointNet实现了点集的排列不变性，并且能够高效地提取每个点的局部特征和整个点云全局特征。在特征编码器中使用的是它提取的全局特征。决策牙齿排列的质量包括：每颗牙齿相对于其他牙齿的位置和方向、每颗牙齿的信息和整个齿列的信息。因此，我们提取颌骨级特征$x_w=E_w(\tilde{X},C)$和牙齿级特征$x_v=E_v(\tilde{X}_v)$，其中$E_w$表示整个牙冠集合的编码器，$E_v$表示单个牙齿的编码器。
特征传播模块
请注意，颌骨级的特征相当稀疏，没有捕捉许多几何细节，如图8(b)所示，细节在Sec4.5中讨论。然而，牙齿级特征捕捉的细节是独立编码的，并且无法关注来自其他牙齿的信息。因此，利用这些特征很难实现牙齿的精确排列。我们引入了基于图的特征传播模块（Feature Propagation Module, FPM），它允许几何细节信息通过图的连接在牙齿之间传递。
图2：(a)牙齿矫正中牙齿排列的代表性例子。每一栏包括治疗前（上）和治疗后（下）的牙齿模型；(b)用于特征传播的牙齿图。这里，我们显示了上颌的牙齿连接，还演示了钳口之间的连接。
我们的特征传播模块$\mathcal{G}$基于[^19]中的传播模型。首先，我们定义了牙齿图$\mathcal{G}=(\mathcal{N,E,H})$，其中$\mathcal{N}$是结点的集合，每个结点对应一个牙齿$v$，$\mathcal{E}$是图的无方向边的集合，$\mathcal{H}$是结点嵌入。而且，两个超级结点用于创建上颌和下颌。这两个超级结点的结点嵌入集合使用零向量初始化。任何其他结点的嵌入$h_v$用其特征$x_v$初始化。如图2(b)所示，$\mathcal{E}$的构成有四种类型的边：$\mathcal{E}_A,\mathcal{E}_S,\mathcal{E}_C,\mathcal{E}_J$，即$\mathcal{E}=\mathcal{E}_A\cup\mathcal{E}_S\cup\mathcal{E}_C\cup\mathcal{E}_J$，其中$\mathcal{E}_A$包含了同一颌中相邻牙齿的关系，$\mathcal{E}S$连接了同一颌中左右对称牙齿；$\mathcal{E}C$由每个牙齿结点与其对应颌的超节点之间的连接组成；$\mathcal{E}J$是两个超级结点之间包含单条边的集合。最后，在$K$次迭代中更新局部特征$x_v$，每次迭代具有固定数量的步骤$T$，如下所示：
$$
m_v^{k,t+1}=\sum{w\in N(v)} A{e{vw}}^k h_w^{k,t}
$$
$$
h_v^{k+1,t+1}=\text{GRU}(h_v^{k,t},m_v^{k,t+1})
$$
其中，$N(v)$表示$v$的相邻节点集，$A_{e_{vw}}^k$是图$\mathcal{E}$中每种类型的边的学习矩阵。在我们的实验中$K$和$T$设置为$3$。
为了进一步改进网络的性能，我们为原有特征增加了一个残差连接。最终更新的牙齿特征$x_v’$是在残差操作后得到。
$$
x_v’=x_v+h_v^{K+1,T+1}
$$
姿态回归器
考虑到临床正畸医生的主观判断或者患者的年龄、性别、脸型等诸多因素对牙齿最佳排列布局的影响，我们生成了一组候选来代替仅仅给出一个结果的方式。最初，提出MoN损失是为了从单个图像恢复成3D时对不确定性建模。受MoN损失的启发[^9]，我们引入了条件加权（Conditional Weighting, CW）方案。
这个方案被设计成允许网络产生多个看似合理的排列，并且仍然能够推荐最合适的那个。为了使CW方案起作用，在姿态回归器中，我们仅仅需要在训练中将随机向量$\xi\in\mathcal{N}(0,I)$附加到输入特征中，其中$\mathcal{N}(0,I)$是零均值单位方差的高斯分布。我们在测试中将$\xi$设置为零向量。CW方案的另一部分在于损失函数（Sec3.4）。所有的特征都被组合起来，并且输入到对应的姿态回归器中去预测6D变换参数：
$$
\Theta_v=\Psi_v(C,x_w,x_v’,\xi)
$$
其中，$\Theta_v$的组成$r_v=(r_v^x,r_v^y,r_v^z)$与$t_v\in\mathcal{R}^3$分别用于旋转的角轴表示和平移。
牙齿组装器
然后，预测的变换参数传递到组装器$\Phi$中去变换、组装和生成最终的输出。这个模块通过指数映射将旋转的角轴表示$r_v\in\mathcal{R}^3$映射回旋转矩阵$R_v\in SO(3)$，这个矩阵是可微的。指数映射的表达式$so(3)\rightarrow SO(3)$通过以下方式将李代数与李群联系在一起。
$$
\exp(r_x)=I_{3\times 3}+\frac{\sin\theta}{\theta}r_x+\frac{1-\cos\theta}{\theta^2}r_x^2
$$
其中，$\theta=|r|_2$是旋转的角度。定义$r=(r^x,r^y,r^z)$为角轴表示的旋转向量，并且伴随了反对称矩阵
$$
r_x=
\begin{bmatrix}
0&-r^z&r^y\
r^z&0&-r^x\
-r^y&r^x&0
\end{bmatrix}
$$
然后，给定牙齿$v$的$r_v$，装配器首先使用等式5将其映射到旋转矩阵$R_v$，接着将变换应用到输入点以获得最终的输出点云
$$
X^*={R_v p_v+c_v+t_v|v\in\mathcal{V},p_v\in\tilde{X}_v}
$$
3.4 损失函数
图4：(a)治疗前模型；(b)对应的治疗后模型；(c)使用从治疗前模型(a)中获得的牙齿形状和对应的治疗后模型(b)获得的牙齿姿态生成的对齐模型，这个被用在我们的重建损失的定义中（见Sec3.4）；(d)将治疗后的模型(b)叠加在对齐模型(c)上，然后可视化它们的差异。
几何重建损失
我们的观察是牙齿在治疗过程中几乎保持刚性，我们的网络特性在输入中保持每个牙齿的形状不变，基于我们的观察和我们的网络特性，我们使用迭代最近点方法来对齐预测结果和基准事实中的每对牙齿（见图4(c)）。然后，对于点$X_v^*$，我们基于刚性对齐结果，通过在基准数据$\bar{X}v$搜索最近点找出他们的对应关系$P{\bar{X}}(X_v^*)$。函数$P_{\bar{X}}(\cdot)$表示这个对应关系搜索过程。为了消除由全局刚性变换引起的损失，并且揭示两个排列之间的内在联系，我们通过最小化下面的能量公式求解了全局刚性变换$\Pi$去对齐预测和基准数据（见图4(d)）：
$$
\arg\min_{\Pi}\sum_{v\in\mathcal{V}}|[X_v^*|1]^T-\Pi[P_{\bar{X}}(X_v^*)|1]^T|2^2
$$
其中，$\text{X}\text{v}^*$是$X_v^*$的坐标矩阵。我们通过正交Procrustes分析解决了上述问题。最后，重建损失计算如下：
$$
L{recon}(X^*,\bar{X})=\sum{v\in\mathcal{V}}|[X_v^*|1]^T-\Pi[P_{\bar{X}}(X_v^*)|1]^T|_S
$$
其中，$|\cdot|S$表示$Smooth{L1}$范数[^12]。
几何空间关系损失
为了强调事实，即良好的排列主要由所有牙齿之间的相互空间关系决定，我们将两个点集$S_1,S_2\subseteq\mathcal{R}^3$之间的几何空间关系定义为：
$$
V_{S_1,S_2}=\bigcup_{i\neq j,1<i,j<2}{x-y^*|y^*=\arg\min_{y\in S_i}|x-y|2^2,x\in S_j}
$$
基于简单的观察，即如果齿列在美学和功能上令人满意，那么两颗牙齿之间的距离不应大于阈值$\sigma$，于是我们通过夹紧所有元素到$[-\sigma,+\sigma]$来计算夹紧后的$V{S_1,S_2}$，并且表示为$V_{S_1,S_2}^C$。我们在所有的实验中根据经验设置$\sigma=5.0$。最后，几何空间关系损失函数的计算如下：
$$
L_{spatial}(X^*,\bar{X})=\sum_{q\in\mathcal{N}}\sum_{e\in\mathcal{P}(q)}
|V_{X_q^*,X_e^*}^C-V_{P_{\bar{X}}(X_q^*),P_{\bar{X}}(X_e^*)}^C|_S
$$
其中，$\mathcal{P}(q)=\text{NBR}(q)\cup\text{OPS}(q)$。函数$\text{NBR}(q)$和$\text{OPS}(q)$分别返回相邻结点和对面的颌。如果节点$q$是表示颌骨的超级节点，那么$X_q^*$表示这个颌骨的所有牙齿的点集。
条件加权损失
如Sec01中所述，我们引入了一种机制，允许网络对不确定性建模，并且为给定的条件向量$\xi$的一个输入生成分布式的输出。我们的方法受到具有以下变化的MoN[^9]损失的启发。我们使得网络能够通过使用条件加权损失来推荐最有可能的安排，条件加权损失的定义如下：
$$
\sum_{\xi_j\sim\mathcal{N}(0,\text{I}),1\leq j\leq n}
\min{\frac1{e^{|\xi_j|}}\cdot Loss(X^*,\bar{X})}
$$
其中，$Loss=L_{recon}+L_{spatial}$，并且在实验中$n=2$。
3.5 实现与训练细节
网络细节
全局的和局部的PointNet编码器编码的特征的维数分别是1024和512.嵌入在FPM中结点的长度设置为512.随机条件向量$\xi$是一个32维向量。姿态回归器由3个线性层组成，其中前两层使用ReLU激活函数和dropout(0.3)，仅在最后一层使用Tanh激活函数。回归器最后一层的权重被初始化为零，因为我们假设牙齿更有可能是不进行移动。
训练细节
在训练开始之前，搜索对应的点对工作就已经完成，因为对应的点对不会由于牙齿的刚性移动而发生改变。在治疗前或者治疗后都没有在模型中出现的牙齿被视为缺失或者拔除。我们在每颗牙齿上随机抽取400个点作为输入。至于缺牙，我们把它们的位置设置为0。为了增加训练的数据量，输入模型的所有单个牙齿，包括治疗前和治疗后的模型，全部随机旋转一个角度，在我们的实验中角度范围在$[-30,+30]$之间，其随机方向和平移距离向量的取值均来自零均值高斯分布$\mathcal{N}(0,I)$。完整的牙齿集合还可以通过随机旋转实现数据增强。请注意，这些增强模型仅被用于扩大模拟的治疗前模型的数据集合。我们假设增强的治疗前牙齿模型$\mathcal{M}^*$应该被映射到对应的治疗后模型$\bar{\mathcal{M}}$，并且增强的治疗后模型$\bar{\mathcal{M}}^*$应该被映射到对应的增强前的原始的治疗后模型$\bar{\mathcal{M}}$。
我们的网络是用PyTorch实现，并且在一个服务器的一个1080-Ti GPU 上训练。我们使用 Adam 优化器。batch size=16, learning rate=1.0e-4，当验证损失停止改善时学习率下降0.5.（ToDo：原文无法理解）
Sec04 实验
4.1 数据集
我们的数据集由300名患者的牙齿模型组成，其中男性（$47%$）和女性（$53%$）的年龄从6岁到18岁不等。对于每个患者，治疗前与治疗后分别扫描得到两个模型。所有三种类型的错位咬合（即类I、II、III）都能够在我们的数据集中观察到，分类标准基于Angel分类[^2]。一些例子如图2(a)所示。对于网络训练，我们对数据集中300对牙齿模型随机地分组为：200对训练集、30对验证集和70对测试集。
4.2 评估指标
我们使用ADD度量[^15]来评估网络预测的精度，ADD度量是预测模型与基准模型逐点计算距离后的平均值。我们还报告了PA-ADD，它是在预测颌和基准颌之间通过Procrustes分析完成刚性配准后计算的ADD。此外，我们定义了PCT@K度量作为网络预测的牙齿误差最小值小于阈值K的百分比。误差可以是形状重建误差、旋转或者平移估计误差等等。类似于6自由度姿态估计中的AUC[^39]，我们定义了PCT-AUC作为PCT曲线下的面积，它是PCT相对于K的积分。形状重建、旋转或者平移误差的PCT-AUC分别被表示为ADD-AUC、$\Delta\theta$-AUC和$\Delta$T-AUC。我们设置PCT-AUC的最大K值