TeethGenerator: A two-stage framework for paired pre- and post-orthodontic 3D dental data generation

We will release the code and dataset of the paper TeethGenerator: A two-stage framework for paired pre- and post-orthodontic 3D dental data generation in this repository.

Dataset

See the Data Use Agreement for details.

Install Requirements

We test our model with Python 3.10 and PyTorch 2.7.1+cu118.

conda create -n teeth_generator python=3.10
conda activate teeth_generator
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install -r requirements.txt

Data Pre-process

Please refer to TADPM or alternatively, you can handle it using your own code. The steps you need to follow are:

• Normalize and center the entire teeth model.
• Extract each individual tooth and use the FPS algorithm to sample 128 points from each tooth.
• Make sure that the sampled points of the teeth before and after orthodontics correspond one-to-one.

After preprocessing, your data should have the following structure:

Path_to_your_pcd_before_orthodontic/
    0_1.ply
    0_2.ply
    ...
    i_j.ply  # Each .ply file corresponds to a single tooth point cloud model before orthodontic treatment.

Path_to_your_pcd_after_orthodontic/
    0_1.ply
    0_2.ply
    ...
    i_j.ply  # Each .ply file corresponds to a single tooth point cloud model after orthodontic treatment.

Stage I

VQ-VAE Training

run the following command:

bash scripts/train_vqvae.sh

Make sure to modify the parameters in scripts/train_vqvae.sh and config/vqvae.yaml.

VQ-VAE Testing

To test or visualize the trained VQ-VAE, run the following command:

bash scripts/test_vqvae.sh

You can modify the parameters in scripts/test_vqvae.sh and adjust the testing logic in runners/test_vqvae.py as needed.

Diffusion Training

run the following command:

bash scripts/train_diffusion.sh

You need to modify the parameters in scripts/train_diffusion.sh and config/diffusion_vae.yaml.

Diffusion Inference

run the following command:

bash scripts/sample_diffusion.sh

You can modify the parameters in scripts/sample_diffusion.sh and adjust the inference logic in runners/test_diffusion.py

Stage II

Training

run the following command:

bash scripts/train_transformer.sh

Inference

run:

bash scripts/sample_transformer.sh

Note that for inference in Stage II, you will need to modify the parameters in config/second_stage_sample.yaml. The style path parameter specifies the data path to the teeth models with the desired style, while the after_path parameter points to the path of synthetic samples generated by Stage I. Your synthetic samples should also follow this structure:

Path_to_your_synthetic_data/
    0_1.ply
    0_2.ply
    ...
    i_j.ply  # Each .ply file corresponds to a single synthetic tooth point cloud model.