PMTSeg: Photo-Guided Tooth Segmentation on 3D Oral Scan Model

Official implementation of CVPR 2026 paper: "Photo-Guided Tooth Segmentation on 3D Oral Scan Model".

• Python >= 3.8
• PyTorch (with CUDA for GPU training)
• CUDA (for PointNet2 CUDA extensions)
• CMake >= 3.27 (for revras)
• OpenGL / EGL (for revras offscreen rendering)
• nvcc (NVIDIA CUDA Compiler)

git clone <repository-url>
cd PMTSeg
git submodule update --init --recursive

pip install -r requirements.txt

PointNet2 provides CUDA operators for point cloud sampling, ball query, etc. Build it first:

cd src/models_3rd/pointnet2
pip install -e .
cd ../../..

If build fails, ensure the CUDA toolkit matches your PyTorch version and nvcc is in your PATH.

revras is used to render id_map, normal_map, and depth_map from meshes. It requires CMake and OpenGL/EGL:

cd revras
pip install -e .
cd ..

Dependencies:

• CMake >= 3.27: Build system
• OpenGL EGL: Headless rendering
  • Ubuntu: sudo apt install libegl1-mesa-dev libgl1-mesa-dev
  • Other systems: Install the corresponding OpenGL/EGL development libraries

revras automatically fetches Eigen, pybind11, etc. via FetchContent.

Ensure the project root is in PYTHONPATH when running scripts (e.g., PYTHONPATH=. python ...) for correct imports of revras and data_preparation.

1. Prepare Teeth3DS or compatible format data as described in the paper.
2. Run the data collection script from the project root:

# Run from project root (PMTSeg)
PYTHONPATH=. python data_preparation/data_collect.py \
  --data-path /path/to/Teeth3DS \
  --out-path work3_front_h5 \
  --img-dir work3_render_2d2 work3_render_2d_front \
  --split-dir Teeth3DS_train_test_split

If using the CAM module for pose estimation, specify the cam output directory with --cam-pose-dir.

cd scripts
python train.py \
  --cfg configs/pmtseg.yml \
  --data-path /path/to/Teeth3DS/work3_front_h5 \
  --model-name model_pmtseg \
  --run-name exp1 \
  --log-dir ./logs \
  --classes 16 \
  --batch-size 4 \
  --epochs 300 \
  --gpu 0

Or use a YAML config file:

python train.py --cfg configs/pmtseg.yml

Run evaluation on the test set with a trained checkpoint:

cd scripts
python test.py \
  --cfg configs/pmtseg.yml \
  --data-path /path/to/Teeth3DS/work3_front_h5 \
  --model-name model_pmtseg \
  --run-name exp1 \
  --log-dir ./logs \
  --classes 16 \
  --gpu 0 \
  --out-dir visualize

This will:

• Load best_1.ckpt from the experiment log directory
• Run inference on the test set
• Save predictions as .ply meshes, .json labels, and .jpg 2D overlays in {log_dir}/visualize
• Compute metrics (TLA, TSA, TIR) and write test_report.csv to the log directory

PMTSeg/
├── src/
│   ├── cam/              # CAM camera alignment module
│   ├── data/             # Dataset (FusionDataset)
│   ├── models/           # Model definitions (PMTSeg)
│   ├── models_3rd/
│   │   └── pointnet2/     # PointNet2 CUDA extensions (requires build)
│   ├── utils/            # Utilities (point_proj, etc.)
│   └── ...
├── revras/               # Reversible Rasterizer (submodule, requires build)
├── data_preparation/     # Data collection scripts
├── scripts/              # Training and testing scripts (train.py, test.py)
├── evaluation/           # Evaluation metrics
└── configs/              # Configuration files

Q: PointNet2 build fails with nvcc not found
A: Install the CUDA Toolkit and add the directory containing nvcc to your PATH.

Q: revras build fails with Could not find OpenGL
A: Install system OpenGL/EGL development libraries, e.g., libegl1-mesa-dev, libgl1-mesa-dev.

Q: from revras import ReversibleRasterizer fails
A: Ensure you have run pip install -e revras/ and that the project root is in PYTHONPATH (e.g., PYTHONPATH=. python your_script.py).

@InProceedings{Zhuang_2026_CVPR,
    author    = {Zhuang, Shaojie and Wei, Guangshun and He, Jiangxin and Zhou, Yuanfeng},
    title     = {Photo-Guided Tooth Segmentation on 3D Oral Scan Model},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    month     = {June},
    year      = {2026},
    pages     = {37558-37567}
}

Please follow the respective licenses of this project and its submodules.