CR-Net: A robust craniofacial registration network by introducing Wasserstein distance constraint and geometric attention mechanism

Abstract

Accurate registration of three-dimensional (3D) craniofacial data is fundamental work for craniofacial reconstruction and analysis. The complex topology and low-quality 3D models make craniofacial registration challenging in the iterative optimization process. In this paper, we proposed a craniofacial registration network (CR-Net) that can automatically learn the registration parameters of the non-rigid thin plate spline (TPS) transformation from the training data sets and perform the required geometric transformations to align craniofacial point clouds. The proposed CR-Net employs an improved point cloud encoder architecture, a specially designed attention mechanism that can perceive the geometric structure of the point cloud. In order to align the source and target data, Wasserstein distance loss is introduced to combined with Chamfer loss and Gaussian Mixture Models (GMM) loss as an unsupervised loss function dedicated to improves registration accuracy. After efficient training, the network can automatically generate the transformation parameters for registration, transforming the reference craniofacial data to the target craniofacial data without manual calibration of feature points or performing an iterative optimization process. Experimental results show that our method has high registration accuracy and is robust to low-quality models.