TADPM：Automatic Tooth Arrangement with Joint Features of Point and Mesh Representations via Diffusion Probabilistic Models

This is the official PyTorch implementation of our CAGD paper "Automatic Tooth Arrangement with Joint Features of Point and Mesh Representations via Diffusion Probabilistic Models."

[paper link]

Installation

First create a conda environment:

conda create --name tadpm
conda activate tadpm

Pytorch / Python combination that was verified to work is:

• Python 3.10, Pytorch 2.3.1, CUDA 11.8

To install python requirements:

pip install -r requirements.txt

To install chamfer distance:

cd chamfer_dist
python setup.py install

To install manifold, please refer to https://github.com/ZhaoHengJiang/MeshReconstruction/tree/main/Manifold

Dataset

See the Data Use Agreement for details.(If you have trouble viewing the DUA on the GitHub page, please clone the repository.)

Data pre-process

• First you need to extract single tooth meshes from dental models, run:

bash scripts/get_mesh.sh

Note that this script automatically centers and normalizes the mesh. You may adjust the normalization scale within the script as needed.

• To get pointcloud files from extracted single tooth meshes:

bash scripts/get_pointcloud.sh

This script extracts corresponding points from individual teeth before and after orthodontic treatment.

• To get remeshed files from extracted single tooth meshes, you can run:

bash scripts/remesh.sh

• To get ground truth rotation parameters between pre- and post-orthodontic dental models , you can run:

bash scripts/register.sh

You need to adjust the data directory parameters in all the scripts mentioned above accordingly.

Pretraining

To pretrain MeshMAE:

bash scripts/pretrain.sh

You can also refer to https://github.com/liang3588/MeshMAE for more details.

You need to modify the file parameter in config/pretrain.yaml, replacing it with the path to a .txt file. Each line of this text file should contain the full path to a remeshed single tooth that will be used for pretraining.

Training

To train the TADPM model:

bash scripts/train.sh

When training TADPM, you should specify the path to the pretrained MeshMAE model checkpoint in this script.

Similarly, you need to modify the parameters in config/TADPM.yaml.

• dataroot should point to the directory containing the remeshed data
• paramroot should specify the directory where the rotation parameters (generated by running scripts/register.sh) are stored.
• before_path and after_path should indicate the directories containing the point cloud files before and after orthodontic treatment, respectively.

Evaluation

Once training is complete, you can run:

bash scripts/test.sh

BibTeX

@article{lei2024automatic,
  title={Automatic tooth arrangement with joint features of point and mesh representations via diffusion probabilistic models},
  author={Lei, Changsong and Xia, Mengfei and Wang, Shaofeng and Liang, Yaqian and Yi, Ran and Wen, Yu-Hui and Liu, Yong-Jin},
  journal={Computer Aided Geometric Design},
  volume={111},
  pages={102293},
  year={2024},
  publisher={Elsevier}
}