TeethSplineSeg: B-Spline Embedded Structure Learning for 3D Tooth Segmentation

Official PyTorch implementation of B-Spline Embedded Structure Learning for 3D Tooth Segmentation

TeethSplineSeg combines a Point Transformer V3 backbone with tooth-centroid-anchored B-spline structural supervision and a Structure-Aware Dynamic Classifier (SADC). The predicted continuous arch coordinate regularizes point features and calibrates sample-specific tooth classifiers using spatial proximity and inter-tooth relations.

Paper: B-Spline Embedded Structure Learning for 3D Tooth Segmentation

---

Main results on 3DTeethSeg22: 96.01% OA, 92.62% T-mIoU, 94.90% Dice, 71.36% B-IoU, 97.50% TIR, and 92.67% TIR=1 (see the paper for details).

---

Main files

• prepare_dataset.py: unified preprocessing from raw OBJ/JSON files to the 32000-point training and testing caches.
• bspline.py: tooth-centroid B-spline fitting and vertex-wise arch-parameter construction.
• models\: implementation of network.
• train_network.py: training entry point for the final full model.
• inference.py: 32000-point prediction, vertex upsampling, post-processing, timing, and evaluation.
• evaluate.py: reports T-mIoU, Dice, OA, B-IoU, TIR, and TIR=1.

Environment

The reference environment is Ubuntu with Python 3.11, PyTorch 2.7.1 + CUDA 12.6, spconv 2.3.8, and FlashAttention 2.8.3. Install the listed dependencies in a compatible CUDA environment:

python -m pip install --upgrade pip setuptools wheel ninja packaging
python -m pip install -r requirements.txt

Depending on the local CUDA toolchain, FlashAttention may need a separate installation using --no-build-isolation.

Data Preparation

Obtain the Teeth3DS dataset from here, and make sure its root contains the official raw OBJ meshes, JSON vertex labels, and split files such as training_lower.txt and testing_upper.txt.

Run the following command from the repository root:

python prepare_dataset.py \
  --source /path/to/Teeth3DS_Dataset

The script links the raw dataset to data/3dteethseg/raw and writes the final cache to:

data/3dteethseg/processed_32000_points_7dim_arch_gsn_jawtype

For every mesh, preprocessing maps upper and lower FDI labels to shared labels 1--16, fits the centroid-anchored B-spline, caches the per-vertex arch parameter beside the raw sample, applies deterministic farthest-point sampling to 32000 points, performs isotropic normalization, and constructs the seven input features [xyz, normal, jaw_type].

If you train this model on other datasets, please make sure that they are converted to the same format as Teeth3DS dataset.

Checkpoints

We provide a checkpoint pretrained on the training set of Teeth3DS in Release. After data preparation step, you can run the inference and evaluation with

python inference.py --ckpt best.ckpt

Training

The default command trains the final model for 100 epochs on one GPU with a batch size of 8:

python train_network.py \
  --experiment_name teethsplineseg \
  --experiment_version full \
  --devices 0 \
  --test_after_train

The defaults use the PTv3 U-Net backbone, predict the B-spline parameter after decoder stage 1, and train SADC with the t-prior gate and tooth-wise prototype attention. Checkpoints and TensorBoard logs are written under logs/. With --test_after_train, the checkpoint with the best validation TIR is passed to inference.py automatically.

Inference

Place a checkpoint at any accessible path, then run the full testing pipeline:

python inference.py \
  --ckpt /path/to/model.ckpt \
  --devices 0 \
  --output_dir results/full_test

Inference produces three prediction stages:

results/full_test/original       # sampled 32000-point labels
results/full_test/upsampled      # labels propagated to original mesh vertices
results/full_test/postprocessed  # final connected-component and graph-cut result
results/full_test/timing.json
results/full_test/postprocessed_summary.txt

To process one testing sample only, use its split name, for example:

python inference.py \
  --ckpt /path/to/model.ckpt \
  --devices 0 \
  --case 014ZWMZ8_lower \
  --output_dir results/014ZWMZ8_lower

If the requested 32000-point cache is missing, this command will preprocesses that case.

Evaluation

To evaluate an existing post-processed prediction directory:

python evaluate.py \
  --pred_dir results/full_test/postprocessed \
  --gt_dir data/3dteethseg/raw \
  --log_file results/full_test/postprocessed_summary.txt

Acknowledgments

This project builds on ideas and open-source implementations from Point Transformer V3, Context-Aware Classifier, 3DTeethSAM, DilatedToothSegNet, and so on. We sincerely thank their authors and contributors for making their work publicly available.

Citation

If you find this repository useful in your research, please cite:

@misc{wei2026bsplineembeddedstructurelearning,
      title={B-Spline Embedded Structure Learning for 3D Tooth Segmentation},
      author={Xianghan Wei and Jianwen Lou and Zhiguo Lu and Hairong Jin and Haihua Zhu},
      year={2026},
      eprint={2608.17291},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2608.17291},
}