JeethuSri/teeth-reconstruction

Tooth Wear Reconstruction using Statistical Shape Models

Reconstruct worn/damaged EDJ (enamel-dentine junction) tooth surfaces from 3D mesh data using PCA-based Statistical Shape Models (SSM) with per-tooth neighborhood-adaptive priors and non-rigid refinement.

Results

Global SSM vs Neighborhood SSM (25 worn teeth)

The neighborhood approach builds a per-tooth local SSM from the nearest anatomically-similar good teeth, instead of one global SSM built from all 15. Averaged across all 25 worn teeth:

Neighborhood wins on 7 of 11 metrics, with the strongest gains on surface-coverage metrics:

SSM trained on 15 unworn ULM3 teeth. Evaluation set: 9 real worn teeth + 8 artificially-worn levels of TEST1 + 8 artificially-worn levels of TEST2 = 25 teeth total.

Adding GPMM and Patch / Hole Filling to the Comparison

The table above compares only the two SSM variants, each reconstruction scored against its own raw worn tooth scan (the only ground truth available for the 9 real worn teeth, which have no known original). The same comparison was extended to the two other reconstruction methods this project now includes, GPMM and Patch / Hole Filling.

A local copy of the raw worn scan, verified point by point to be free of read corruption, was only available for 19 of the original 25 teeth: all 16 TEST1 and TEST2 levels, plus real teeth 06, 07, and 08. Real teeth 01, 02, 03, 04, 05, and 09 were left out; their local copies each had thousands of corrupted vertices (or, for tooth 03, a smaller but still meaningful number) that would otherwise silently distort these numbers, especially Hausdorff and R squared. All four methods below are scored on this same 19 tooth set, so the comparison across the four columns is fair. As a sanity check, the Global and Neighborhood SSM numbers recomputed here land close to the 25 tooth table above (R squared 99.89 percent here versus 99.74 percent there, Chamfer 0.070 mm here versus 0.088 mm there), which is the expected amount of difference for a slightly smaller, verified subset of teeth against a separately sourced copy of the same raw scans.

Mean across the 19 teeth, reconstruction scored against the raw worn tooth scan:

Best method per metric, count out of the 19 teeth:

On this metric, reconstruction scored against the raw worn scan rather than the true original, the two SSM variants dominate. This is expected rather than a sign GPMM and Patch / Hole Filling are worse restorers: a method that changed nothing at all from the worn input would score perfectly here, so this metric rewards staying close to the observed worn surface more than it rewards recovering the true, unworn anatomy. GPMM in particular refits the whole surface under its shape prior, so it moves further from the raw worn scan than the SSM methods even when it is recovering worn anatomy well. The evaluation against the true original tooth, further down this README and in full detail in results.md, is the fairer test of restoration quality and tells a different story.

Full per tooth numbers behind this table: `ssm_pipeline/output/eval_flagship_19teeth.csv`.

Reconstruction Methods

Since the Global versus Neighborhood SSM comparison above, the pipeline has grown three additional reconstruction strategies. All four are evaluated head to head on two independent datasets.

Results Summary

Every method was scored against the true original (unworn) tooth on two independent datasets: the original 16 case set (TEST1 vs n0245, TEST2 vs n0257) and a newer 64 case set (8 real specimens x 8 Molnar wear levels). Two metrics were used per case: a full surface metric (every point) and a worn region metric (restricted to points genuinely worn away, the fairer restoration test). Full detail, every case and every metric, is in results.md.

Old dataset, mean of 16 cases. Plain SSM (Global and Neighborhood, nearly tied) is the best restorer once scored on the worn region only. Patch/Holes has by far the best full surface score because it changes the fewest points of any method, but that is also why it is the worst restorer of the four once judged fairly on what it actually filled in.

v5 dataset, mean of 64 cases. Neighborhood SSM and Global SSM are essentially tied for best among the three index paired methods. Patch/Holes shows the lowest numbers overall, but its worn region there is self selected and smaller on average than the shared mask used by the other two, so this particular comparison is not perfectly fair. See results.md for the full explanation.

Testing Status

Done:

• —All 4 methods (Global SSM, Neighborhood SSM, Patch/Holes, GPMM) evaluated on the old dataset (16 cases) with both metrics.
• —All 4 methods evaluated on the v5 dataset (64 cases) with both metrics.
• —Correspondence and reconstruction pipeline re run on a new, larger training set (90 good teeth, 8 held out test specimens, 10k points per tooth, --max-neighbors 10) to check the approach generalizes beyond the original 15 tooth training set.

Not yet done:

• —Reconcile the worn region mask methodology between the old and v5 datasets so Patch/Hole Filling's v5 result is directly comparable to the other methods.
• —Investigate the severe wear (Molnar stage 6, wear level 7) failure mode where Patch/Holes degrades sharply relative to plain SSM.
• —Seam quality polish on Patch/Holes meshes (ragged re opened cervical edge, occasional proud bump at a filled hole).

Full numeric detail for everything above: results.md.

Data Analysis — Shape Space Visualizations

These figures are produced by data_analysis/all_teeth_analysis.py after Stage 2, and give a geometric view of how the 15 good teeth, 25 worn teeth, and their reconstructions distribute in shape space.

PCA Scree Plot — How Much Variance Each Mode Captures

[image]

The first two PCA modes alone capture ~75 % of total variance across the 15 good teeth, and ~95 % is reached by mode 8. This is why a 2D PCA plot is a faithful summary of tooth shape — and why local SSMs (K ≤ 5 neighbors) only need 1–4 modes.

PCA 2D — Good Teeth, Worn Teeth, and Reconstructions

[image]

Each tooth is one point in the PCA shape space of the 15 good teeth. Blue circles = good teeth, orange = originals, red triangles = real worn inputs, green diamonds = TEST1 wear levels, purple squares = TEST2 wear levels. Stars and X-markers overlay the corresponding reconstructions.

Two things jump out:

• —TEST1 and TEST2 form tight per-tooth clusters — progressive wear barely moves a tooth in PCA space, which is why reconstructions stay close to the input.
• —Real worn teeth 03, 06, 07 sit in the upper-left corner far from the main cluster — this is exactly why the neighborhood algorithm selects only tooth_14 (K=1) for these: they are outliers and averaging across all 15 good teeth would pull the reconstruction back toward a centroid they don't belong to.

PCA 2D — Zoomed View (outlier cluster removed)

[image]

Same PCA space, but with the bottom-right outlier cluster (T03, T05, T06, T07 worn and T14 good) excluded so the main population is easier to read. This view shows both reconstruction families on the same axes — Recon (Global) in purple X markers (built from the global SSM over all 18 good teeth) and Recon (Nbr) in pink stars (built from each worn tooth's local SSM of up to 5 nearest neighbors). The two reconstructions land in clearly different places for the same input, which is the visual signature of how much the neighborhood prior pulls each tooth toward its local subgroup instead of the global centroid.

t-SNE (Raw 300,000-D features) — Non-linear Shape Embedding

[image]

t-SNE on the full 300,000-D point-cloud features (before PCA) shows the same neighborhood structure that PCA exposes, but non-linearly. TEST2 forms a very tight cluster in the lower-right — every wear level of TEST2 shares nearly identical overall shape. TEST1 fans out along a roughly linear "wear trajectory" on the left side. Real worn teeth are scattered among the good teeth, confirming each one has its own anatomically-similar subset — the core motivation for the neighborhood approach.

Worn-to-Reconstruction Distance in PCA Space

[image]

Euclidean distance in PC1–PC2 space between each worn tooth and its reconstruction. Small bars mean the reconstruction lands near its input — which is what we want. T07, T06, T03 show the largest displacements: these are exactly the teeth the neighborhood algorithm flagged as outliers (K=1 with tooth_14 as sole neighbor), because the SSM has to pull a heavily-worn tooth toward a complete-anatomy prior. The TEST1 and TEST2 reconstructions cluster near zero, confirming the pipeline is stable on clean inputs.

A full gallery of paired-distance plots (t-SNE, UMAP, local variants) is in `data_analysis/plots_v2/`.

What This Project Does

Teeth wear down over time through mastication, bruxism, dietary abrasion, and chemical erosion. This wear removes original cusp geometry from the enamel-dentine junction (EDJ), making morphological analysis difficult in dental anthropology and paleontology.

This pipeline learns the statistical shape variation of unworn upper-left third molars (ULM3) from 15 complete specimens, then reconstructs missing anatomy on worn teeth in two stages:

1. 1.Global SSM reconstruction — one PCA model built from all 15 good teeth, used as baseline.
2. 2.Neighborhood SSM reconstruction — for each worn tooth, the pipeline automatically selects its anatomically-nearest good teeth and builds a local SSM tailored to that tooth's shape family. This gives a tighter shape prior and measurably better surface coverage.

Output per worn tooth: a 100k-point reconstructed point cloud, a smooth watertight triangle mesh, the reconstruction in the original millimeter input space, and a full evaluation JSON.

Pipeline Overview

                        TOOTH RECONSTRUCTION PIPELINE

  +-------------------+        +-------------------+
  |  Good Teeth       |        |  Worn Teeth       |
  |  Unworn ULM3 PLY  |        |  Real + Test PLY  |
  +---------+---------+        +---------+---------+
            |                            |
            v                            v
  +----------------------------------------------------------+
  |  STAGE 1: CORRESPONDENCE                                 |
  |  correspondence_pipeline.py                              |
  |                                                          |
  |   1. Sample uniform points per tooth                     |
  |   2. Normalize (center, scale to bbox diag, PCA align)   |
  |   3. ICP rigid alignment to auto selected template       |
  |   4. Coarse to fine CPD non-rigid registration           |
  |      (CPD on a subset, then KNN upsample to full count)  |
  |                                                          |
  |  Output: corresponded.ply per tooth                      |
  |          (shared point ordering across every tooth)      |
  +-------------------------+--------------------------------+
                            |
          +----------+---------+----------+
          |          |         |          |
          v          v         v          v
     +--------+ +--------+ +--------+ +--------+
     |STAGE 2a| |STAGE 2b| |STAGE 2c| |STAGE 2d|
     |GLOBAL  | |NEIGHBOR| |PATCH/  | |GPMM    |
     |SSM     | |HOOD SSM| |HOLES   | |        |
     |        | |        | |        | |        |
     |reconst-| |neighbor| |patch_  | |gpmm_   |
     |ruction_| |hood_   | |reconst-| |reconst-|
     |pipe-   | |reconst-| |ruction.| |ruction.|
     |line.py | |ruction.| |py      | |py      |
     |        | |py      | |(holes  | |        |
     |        | |        | |mode)   | |        |
     +---+----+ +---+----+ +---+----+ +---+----+
         |          |          |          |
         +----------+----------+----------+
                             |
                             v
  +----------------------------------------------------------+
  |  STAGE 3: EVALUATION AND ANALYSIS                        |
  |  evaluate_all_methods.py, data_analysis/                 |
  |                                                          |
  |  Full surface and worn region Chamfer, RMSE, Hausdorff,  |
  |  coverage; PCA, t-SNE, and UMAP shape space plots        |
  +----------------------------------------------------------+

Stage 1 produces one shared, point aligned representation of every tooth. Stages 2a through 2d are four independent reconstruction methods that all read that same Stage 1 output; none of them depend on each other, so any subset can be run on its own. Stage 3 scores whichever methods were run, against the true original tooth where one is available.

How Neighborhood Selection Works (4 Steps)

For every worn tooth, the pipeline decides which good teeth to build the local SSM from:

1. 1.PCA on good teeth. All 15 good teeth are flattened and projected into a low-dimensional PCA shape space (up to 10 components, 95% variance). Each tooth becomes one point in that space.
2. 2.Adaptive distance threshold. All pairwise distances between the 15 good teeth are computed in PCA space. The threshold is set at median + 2 × IQR of those pairwise distances (≈ 13.72 for this dataset). Anything farther is considered anatomically "too different."
3. 3.Nearest-neighbors within threshold. The worn tooth is projected into the same PCA space. Distances to all 15 good teeth are computed, and every tooth closer than the threshold is selected — capped at 5.
4. 4.No minimum required. Even a single neighbor is used (that tooth's shape becomes the local mean, 0 PCA modes). Only if zero good teeth are within the threshold does the method fall back to the global SSM. In practice this fallback triggered on only 1 of 25 worn teeth (tooth_05, whose nearest good tooth sat at distance 30 — well past the 13.72 threshold).

Example neighbor assignments from this dataset:

• —tooth_01 → neighbors [08, 10, 04, 15, 11] (K=5)
• —tooth_03, tooth_06, tooth_07 → neighbor [14] only (K=1, local mean = tooth_14)
• —tooth_05 → global fallback (no neighbor within threshold)
• —tooth_TEST2_level7 → neighbors [09, 03, 05, 12, 06] (K=5)

Project Structure

Teeth-Reconstruction-/
├── all_good_teeth/                 # 15 unworn ULM3 EDJ meshes (original SSM training set)
│   └── cprc_nyu_*.ply
├── all_worn_input/                 # 25 worn tooth inputs (9 real + TEST1 + TEST2)
│   ├── tooth_01/wear_real.ply      # Real worn teeth
│   │   ...
│   ├── tooth_09/wear_real.ply
│   ├── tooth_TEST1/
│   │   ├── wear_level0.ply         # Artificially-worn (8 levels each)
│   │   │   ...
│   │   └── wear_level7.ply
│   └── tooth_TEST2/
│       └── wear_level0..7.ply
├── new_good teeth_ULM3's/          # 98 unworn ULM3 teeth (larger v5 training pool)
├── Artificially Worn/               # 8 real specimens x 8 Molnar wear levels (v5 test set)
│
├── ssm_pipeline/
│   ├── correspondence_pipeline.py       # STAGE 1: point correspondence
│   ├── reconstruction_pipeline.py       # STAGE 2a: global SSM reconstruction
│   ├── neighborhood_reconstruction.py   # STAGE 2b: neighborhood-adaptive SSM
│   ├── patch_reconstruction.py          # patch / hole-filling reconstruction (--detect-mode holes)
│   ├── gpmm_reconstruction.py           # GPMM posterior shape-completion reconstruction (Tier 1)
│   ├── evaluate_all_methods.py          # consolidated Chamfer/RMSE evaluation, all methods x both datasets
│   ├── local_mean_reconstruction.py     # (legacy) single-tooth local mean
│   ├── correspond_originals.py          # utility: correspond original-space outputs
│   ├── run_correspondence_a100_short.slurm
│   ├── run_full_pipeline.slurm          # correspondence + global recon
│   ├── run_neighborhood_recon.slurm     # neighborhood reconstruction only
│   ├── input_v5/                        # staged v5 inputs (90 good teeth, 8 test teeth reorganized)
│   └── output/
│       ├── correspondence_all_100k/     # Stage 1 output, old dataset (15 good + 25 worn corresponded)
│       ├── correspondence_v5_10k/       # Stage 1 output, v5 dataset (90 good + 64 worn corresponded, 10k pts/tooth)
│       ├── recon_all/, recon_all_v5/               # Stage 2a global reconstructions
│       ├── recon_neighborhood/, recon_neighborhood_v5/  # Stage 2b neighborhood reconstructions
│       │   ├── ssm/
│       │   ├── neighbor_selection.json  # which neighbors each worn tooth used
│       │   ├── comparison.json          # global vs neighborhood per tooth
│       │   └── reconstructions/tooth_XX/
│       │       ├── worn_input.ply
│       │       ├── reconstructed.ply
│       │       ├── reconstructed_in_input_space.ply
│       │       ├── reconstructed_smooth.ply
│       │       ├── coefficients.npy, removed_mask.npy
│       │       └── evaluation.json
│       ├── recon_holes_final/, recon_holes_v5/     # patch/hole-filling reconstructions
│       ├── recon_gpmm_test/, recon_gpmm_v5/        # GPMM reconstructions
│       ├── eval_old_dataset.csv, eval_v5_dataset.csv  # consolidated evaluation results (current)
│       ├── eval_flagship_19teeth.csv    # 4-method comparison vs raw worn tooth, 19 of 25 old-dataset teeth
│       └── archive/                     # superseded experimental iterations (gitignored, local only)
│
├── data_analysis/
│   ├── all_teeth_analysis.py       # PCA + t-SNE + UMAP + comparison plots across all recon outputs
│   ├── good_teeth_pca.py           # PCA on just the good-tooth training set
│   ├── worn_teeth_projection.py    # project worn teeth into good-teeth PCA space
│   └── plots_v2/                   # generated figures
│
├── requirements.txt                # pip freeze of `teeth` conda env
├── worn.png / reconstruction.png   # example visualization (input vs output)
└── README.md

Installation

Requires Python 3.9–3.12 (Open3D does not yet support 3.13+).

Option A — Clone + pip install (matches the cluster environment exactly)

git clone https://github.com/yourusername/Teeth-Reconstruction.git
cd Teeth-Reconstruction
pip install -r requirements.txt

Option B — HPC cluster with conda (NYU-Langone setup)

source /gpfs/data/davolilab/software/conda-envs/miniconda3/etc/profile.d/conda.sh
conda activate teeth

All SLURM scripts in ssm_pipeline/ assume this teeth environment.

GPU support

CPD non-rigid registration and SSM SVD can be GPU-accelerated via CuPy. The pinned requirements.txt installs cupy-cuda12x. For a CUDA 11 system, replace with cupy-cuda11x.

Quick Start

All reconstruction commands assume you are in ssm_pipeline/ with the environment active.

cd ssm_pipeline
conda activate teeth    # or: source venv/bin/activate

Stage 1 — Correspondence (GPU, ~1–4 hours for 40 teeth)

Establishes point-to-point anatomical correspondence across all 15 good + 25 worn teeth.

python correspondence_pipeline.py \
  --good-teeth "../all_good_teeth" \
  --artificial-wear "../all_worn_input" \
  --output "output/correspondence_all_100k" \
  --n-points 100000 \
  --registration-mode coarse2fine \
  --cpd-points 43000 \
  --displacement-knn 3 \
  --auto-template

On the HPC cluster:

sbatch run_correspondence_a100_short.slurm

Stage 2a — Global SSM Reconstruction (CPU, ~10 min for 25 teeth)

Baseline: one SSM built from all 15 good teeth, applied to every worn tooth.

python reconstruction_pipeline.py \
  --correspondence-dir "output/correspondence_all_100k" \
  --artificial-wear "../all_worn_input" \
  --output "output/recon_all" \
  --skip-eval \
  --proxy-missing-fraction 0.25 \
  --variance-threshold 0.99

Or combined with Stage 1 via SLURM:

sbatch run_full_pipeline.slurm

Stage 2b — Neighborhood SSM Reconstruction (GPU recommended, ~20 min)

Per-worn-tooth local SSM from adaptively-selected nearest good-tooth neighbors.

python neighborhood_reconstruction.py \
  --correspondence-dir "output/correspondence_all_100k" \
  --global-recon-dir "output/recon_all" \
  --output "output/recon_neighborhood" \
  --artificial-wear "../all_worn_input" \
  --max-neighbors 5 \
  --threshold-iqr-mult 2.0 \
  --proxy-missing-fraction 0.15 \
  --ssm-variance 0.95 \
  --regularization 1.0

On the HPC cluster:

sbatch run_neighborhood_recon.slurm

Outputs land in output/recon_neighborhood/. Two JSONs summarize the run:

• —neighbor_selection.json — which good teeth each worn tooth used, distances, and fallback reasons
• —comparison.json — side-by-side metrics vs the global reconstruction

Stage 3 — Cross-method Analysis

cd ../data_analysis
python all_teeth_analysis.py \
  --correspondence-dir ../ssm_pipeline/output/correspondence_all_100k \
  --extra-recon-dir ../ssm_pipeline/output/recon_neighborhood/reconstructions

Produces PCA plots, pairwise-distance matrices, and comparison figures in data_analysis/plots_v2/.

Parameters Reference

correspondence_pipeline.py

reconstruction_pipeline.py (Global SSM)

neighborhood_reconstruction.py (Neighborhood SSM)

Output Files

Each worn-tooth directory (output/<run>/reconstructions/tooth_XX/) contains:

The SSM itself is saved to output/<run>/ssm/:

Neighborhood run also produces two top-level JSONs:

• —neighbor_selection.json — per-tooth neighbor assignments, distances, threshold info
• —comparison.json — per-tooth global vs neighborhood metric comparison

Evaluation Metrics

All metrics compare the reconstruction (in original input space) against the raw worn tooth.

Key Algorithms

Statistical Shape Model (SSM)

PCA learns the principal modes of shape variation:

Shape(b) = μ + V · b

μ is the mean shape (100k × 3 flattened to 300k), V the eigenvectors, b the fitted coefficients. Coefficients are fitted to observed (non-missing) points via Tikhonov-regularized least squares:

b = (VᵀV + λΛ⁻¹)⁻¹ Vᵀ (x_obs - μ_obs)

Coefficients are clipped to ±4σ to prevent extreme extrapolation.

Neighborhood SSM

The same PCA math, but μ and V are rebuilt per worn tooth from only its K selected neighbors (K ≤ 5). For K = 1, μ is that single neighbor's shape and V is empty (pure mean-shape fit). Built local SSMs are cached by frozenset(neighbor_indices) so teeth with identical neighbor sets don't recompute the SVD.

Non-Rigid Refinement

After SSM fitting:

1. 1.Observed points are replaced exactly with the worn input coordinates (observation error → 0).
2. 2.Missing points are displaced by KNN-interpolated corrections from nearby observed points, producing smooth transitions.

Screened Poisson Surface Reconstruction

1. 1.Estimate + orient surface normals (outward from centroid)
2. 2.Open3D Screened Poisson (octree depth 9) → watertight triangle mesh
3. 3.Trim bottom 1% of Poisson density (removes extrapolated fringe)
4. 4.30 iterations Taubin smoothing (λ=0.5, μ=-0.53)
5. 5.Fix normals and triangle winding via trimesh

Tooth-to-Specimen Mapping

Real worn teeth (all_worn_input/tooth_01..09/wear_real.ply):

Artificially-worn test teeth: tooth_TEST1/wear_level0..7.ply and tooth_TEST2/wear_level0..7.ply — 8 progressively-worn versions of two source teeth, used to test reconstruction robustness at increasing wear levels.

Good teeth (SSM training set, all_good_teeth/): 15 unworn ULM3 EDJ meshes including n0043, n0047, n0049, n0256, n0258, n0264, n0266, n0269, n0291, n0292, n0293, n0298, n0300, n0307, n0350.

Dependencies

Full pinned list: `requirements.txt`. Headlines:

• —Python 3.9–3.12
• —trimesh >= 4.0.0
• —open3d >= 0.18.0
• —numpy, scipy, scikit-learn
• —probreg >= 0.3.0 (GPU-accelerated CPD)
• —pycpd >= 2.0.0 (CPU fallback)
• —cupy-cuda12x (optional, GPU linear algebra)

Citation

@software{tooth_reconstruction_ssm,
  title  = {Tooth Wear Reconstruction using Statistical Shape Models
            with Neighborhood-Adaptive Priors},
  author = {Jeevan Ananth},
  year   = {2026},
  url    = {https://github.com/yourusername/Teeth-Reconstruction}
}

License

MIT License — see LICENSE for details.