A Self-Supervised Deep Framework for Reference Bony Shape Estimation in Orthognathic Surgical Planning

The simulator network consists of a series of encoding, fusion, and decoding layers to learn point features from the coordinates of vertices (Fig. 3). There are two encoding branches with shared weights to extract point features from the two input bony surfaces separately. Each encoding layer selects a subset of $N_{\text{sub}}$ points from the input points via furthest point sampling [5]. For each sampled point, its neighboring points in a 3D ball region with radius $r$ are gathered from the input points. The point convolution is then performed over point features of the gathered points via PointConv [6], which is implemented based on shared multi-layer perceptron (s-MLP) and max-pooling [7]. Following the encoding layers, the outputs of corresponding layers from two encoding branches are fused via concatenation and s-MLPs. The fused features are then decoded by upsampling via point interpolation [5], grouping via 3D ball neighboring, and convolution via PointConv. The decoding layer increases the number $N_{\text{sub}}$ of points and reduces the dimension $C_{\text{feature}}$ of each feature vector. A set of cascaded fusion and decoding layers are applied to generate a $N\times 3$ output vector $V_{\text{simulator}}$, which updates the positions of the vertices of the input normal bone to obtain the simulated deformed bone.

We train the simulator network using loss function

which encourages the jaw of the simulated deformed bone to be similar to the input deformed bone based on the relative coordinates of vertices. Specially, a set of landmarks are defined on each bony surface, the relative coordinate vector $r_{\text{def\_jaw}}(i,j)$ of the $i$-th jaw vertex relative to the $j$-th landmark on the deformed bony surface is calculated as

where $c_{\text{def\_jaw}}(i)$ and $c_{\text{def\_landmark}}(j)$ are respectively the coordinate vectors of the $i$-th vertex and the $j$-th landmark. Similarly, the relative coordinate vector for the $i$-th jaw vertex of the simulated deformed bony surface is

We calculate $L_{\text{jaw}}$ by

where $N_{\text{jaw}}$ is the number of vertices that belong to the jaw, $K$ is the number of landmarks, and $\left\|\cdot\right\|_{2}$ is the $\ell_{2}$-norm. $L_{\text{midface}}$ forces the network to fix the midface during simulation. To keep the midface unchanged, i.e., $L_{\text{midface}}=0$, we set the displacement vectors corresponding to midface vertices to zero in $V_{\text{simulator}}$. A smooth deformation field is encouraged by calculating

where $u(i)=c_{\text{simu}}(i)-c_{\text{norm}}(i)$ and $u(j)=c_{\text{simu}}(j)-c_{\text{norm}}(j)$ are respectively the displacement vectors of $v_{i}$ and $v_{j}$ on the simulated deformed bony surface. $\mathcal{N}(v_{i})$ is the set of one-ring neighboring vertices of $v_{i}$. $N$ is the total number of vertices on the surface. Finally, a $\ell_{2}$-norm regularization $L_{\text{simu\_reg}}$ of the network parameters is incorporated to avoid overfitting.

The corrector is an encoding-decoding network (Fig. 4). Like the simulator network, each encoding layer in the corrector performs sampling, grouping, and convolution operations; each decoding layer performs interpolation, grouping, and convolution. The skip connection between two corresponding layers from the encoding and decoding streams is applied to learn point features integrating local and global shape information. Following the decoding stream, the point features are processed by a set of s-MLPs to obtain a $N\times 3$ vector, which corrects the positions of vertices of the simulated deformed bone to get a corrected bone.

The loss function for the corrector network is

where $c_{\text{correct}}(i)$ and $c_{\text{norm}}(i)$ are the coordinate vectors of the $i$-th vertex on the corrected and input normal bony surface, respectively, $L_{\text{correct\_reg}}$ enforces $\ell_{2}$-norm regularization on the network parameters.

To train the proposed networks, vertex-wise correspondences of surfaces are established by non-rigidly matching a template surface with all bony surfaces in the training set. First, a group of corresponding landmarks localized on the training surfaces are rigidly aligned. From the aligned surfaces, we then choose the surface that is closest to the average landmarks as the template, and warp it to each rest surface via non-rigid coherent point drift (CPD) matching [8]. All warped surface templates are used for network training. We reduce the number of vertices via surface simplification [9] to reduce the computation cost. The vertex coordinates are min-max normalized. The loss functions are minimized with the Adam [10] optimizer to determine the optimal network parameters. The two networks are alternatively trained in turn on each batch of samples. During the inference stage, the trained corrector is directly applied on a patient bony surface to estimate a reference bony shape model. The corrector network is invariant to the number and the order of vertices on the testing surface [7].

We used CT scans of 67 normal subjects from a previous study [11] and CT scans of 116 patients with jaw deformities to train our networks. The study was approved by our Institutional Review Board (#Pro00009723). Following the clinical routine [12], all CT data were segmented to reconstruct 3D bony shape models. 51 landmarks were localized on each bony surface by experienced oral surgeons for calculating the loss function during the training only. Following Section 2.0.3, a bony surface template from the training surfaces was selected based on the 51 landmarks to establish dense vertex correspondences. Ultimately, a total of 7772 (67 $\times$ 116) random pairs of deformed-normal bones were used to train the networks.

The simulator network was implemented with 4 encoding, 4 fusion, and 4 decoding layers. The corrector network was configured with 4 layers each for encoding and decoding. The numbers of points for the encoding and decoding layers were $N_{\text{sub}}=\{\frac{N}{4},\frac{N}{16},\frac{N}{64},\frac{N}{128},\frac{N}{64},\frac{N}{16},\frac{N}{4},N\}$, where $N=4724$ during training, and $N>10\,000$ during testing. The radius $r$ of the 3D neighboring ball was set to $\{0.1,0.2,0.4,0.8\}$ and $\{0.8,0.4,0.2,0.1\}$ for the encoding and decoding layers, respectively. The output point features from the 4 encoding and 4 decoding layers had dimensions $C_{\text{feature}}=\{64,128,256,512\}$ and $C_{\text{feature}}=\{512,256,128,128\}$, respectively. Empirically, we set $\alpha=0.3$, $\beta=0.1$, and $\lambda=0.1$ in the loss functions (1) and (6). The number $K=51$ of landmarks was set for calculating $L_{\text{jaw}}$ in (4). The networks were trained for 400 epochs with the learning rate of $0.0001$.

For testing, we acquired the data from another 24 patients with paired pre- and post-operative CT scans. The post-operative bones were used as ground truth for performance evaluation. Following the method in [4], we remeshed each post-operative bony surface according to its pre-operative bony surface to construct the vertex-wise correspondences for calculating estimation accuracy. Four evaluation metrics [4], i.e., vertex distance (VD), edge-length distance (ED), surface coverage (SC), and landmark distance (LD), were employed in the quantitative evaluation. We compared our method with DefNet [4], which was trained using 6834 paired samples simulated by sparse representation, 102 deformed bones were synthesized for each of 67 normal bones from 116 patient bones. Qualitative evaluation was performed by an experienced oral surgeon. Statistical significance of accuracy improvement was evaluated using the paired t-test.

The trained framework was tested on the data of 24 patients. Figure 5 shows the estimated bony surfaces and the vertex-wise distance heatmaps for three randomly selected patients¹¹1A video illustration of the three estimation examples is available as a supporting material.. All results from our method inspected by the expert suggest that the estimated bones are clinically acceptable. For quantitative evaluation, the estimation accuracy for the jaw and the midface were calculated separately. Table 1 shows that our method is statistically significantly more accurate in estimating the normal jaw ($p<0.05$) than DefNet based on the four metrics. The accuracy in maintaining the midface is comparable between the two methods ($p>0.05$). Overall, the proposed method outperforms DefNet in the term of accuracy.

Our method takes about 30 minutes for one training epoch, and 10 seconds for testing on each patient data, evaluated based on an NVIDIA Titan Xp GPU.