Skip to content
cvlab-epflPublic

About

repository for PrIntMesh (Precise Intersection Surfaces for 3D Organ Mesh Reconstruction)

Resources

Stars

2 stars

Watchers

0 watching

Forks

Latest commit

 

History

1 Commit

Folders and files

Repository files navigation

PrIntMesh

Precise Intersection Surfaces for 3D Organ Mesh Reconstruction.

PrIntMesh reconstructs anatomically accurate surface meshes of multiple organs directly from a volumetric medical scan. A segmentation backbone predicts voxel labels, and a graph-convolutional mesh decoder deforms a subdivided template into the organ surfaces in a single forward pass — sharing a template so that adjacent organs meet along precise, consistent intersection surfaces rather than interpenetrating or leaving gaps. Training is end-to-end with Chamfer distance plus mesh-regularization losses (edge length, Laplacian smoothing, normal consistency).

The mesh decoder is backbone-agnostic: it plugs on top of any voxel segmentation network. This repository demonstrates that with two backbones:

  • Reference path — heart (MMWHS). Our own U-Net backbone and our own preprocessing. This path is fully self-contained.
  • nnU-Net backbone — hippocampus, lung. The same mesh decoder on top of nnU-Net as the segmenter.

Repository layout

Path Contents
printmesh/ The library (installed as an editable package). Model (model/: lit.py LitUNet, mesh decoders, graph conv, U-Net backbones, losses, metrics, CUDA rasterizer), data (data/), utilities (util/), mesh smoothing (smoothing/).
scripts/ Runnable entry points — train.py (Lightning CLI: fit/test) and preprocess.py (dataset generation).
confs/ One config per released experiment: heart_split{0..4}.yaml (5-fold heart CV), hippo.yaml, lung.yaml.
tools/ Auxiliary helper scripts (dataset json, plotting, conversions).
templates/ Template meshes deformed by the decoder.
pyproject.toml, pixi.toml Packaging and the pinned GPU environment.

Installation

A CUDA GPU is required. The environment is managed with pixi and pinned to torch 2.6 / cu126, kaolin 0.18, and PyTorch3D 0.7.8:

pixi install
pixi run pytorch3d-install   # builds PyTorch3D from source against the installed torch

Quickstart: the heart reference path

The heart path runs end-to-end from raw data with nothing external.

1. Get the raw data. We use the CT training set of the MICCAI 2017 MM-WHS challenge (Multi-Modality Whole Heart Segmentation). After unzipping you should have a ct_train/ directory of 20 paired volumes:

ct_train/
├── ct_train_1001_image.nii    # CT scan
├── ct_train_1001_label.nii    # segmentation (label values 500/600/420/550/205/820/850)
├── ct_train_1002_image.nii
├── ct_train_1002_label.nii
└── ...                        # 1001 … 1020

2. Preprocess into the preprocessed/mmwhs-128 dataset with a single command:

pixi run python scripts/preprocess.py mmwhs heart-centered \
    -f /path/to/MMWHS/ct_train --relabel --meshes -o preprocessed/mmwhs-128 -d cuda

3. Train:

pixi run python scripts/train.py fit --config confs/heart_split0.yaml

heart_split{0..4}.yaml are the five cross-validation folds. Each loads a pretrained U-Net (ckps/unet/split{i}.ckpt) and deforms the template into the four chamber surfaces plus myocardium. Logs and checkpoints are written under tb_logs/.

Evaluate a finished run:

pixi run python scripts/train.py test --config <run_dir>/config.yaml \
    --ckpt_path <run_dir>/checkpoints/best.ckpt

Evaluation writes three files under the run directory:

  • metrics.csv — one row per test sample, metric, and class. dataset_idx identifies the img_NNN input, while sample_idx is its position in the test pass. class_axis is seg, face, or scalar; class_id is blank for scalar hole metrics.
  • summary.csv — per-class and overall means. n_valid is the number of contributing samples for class rows and the number of contributing class means for overall rows.
  • results.txt — a human-readable rendering of the same results.

Missing values are blank in the CSV value column and explained by status: missing_prediction, missing_ground_truth, or non_watertight. Summaries use only rows whose status is ok.

Additional datasets via the nnU-Net backbone

hippo.yaml and lung.yaml set model.unet_name: nnunetd, which loads an nnU-Net plan (config_path) and fold weights (weight_fold) as the segmentation backbone. Preparing these datasets uses nnU-Net's own preprocessing rather than scripts/preprocess.py:

  1. Obtain the raw data — hippocampus is Medical Segmentation Decathlon, lung is from the TotalSegmentator dataset.
  2. Run nnU-Net's standard plan_and_preprocess for the corresponding dataset id, then convert its preprocessed volumes into the img_NNN_{data.npy,seg.npz,info.yaml} format the data loader expects.
  3. Point data.path and the config_path / weight_fold in the config at your local nnU-Net paths (the shipped configs use machine-specific absolute paths that must be updated).

Data availability

We cannot redistribute the datasets. The heart path is reproducible from the public MMWHS scans with the command above; hippocampus and lung follow the nnU-Net preprocessing route described above.

Citation

This code accompanies our paper, PrIntMesh: Precise Intersection Surfaces for 3D Organ Mesh Reconstruction (arXiv). If you use this code or find our work useful, please cite:

@article{mercadier2025printmesh,
  title={PrIntMesh: Precise Intersection Surfaces for 3D Organ Mesh Reconstruction},
  author={Mercadier, Deniz Sayin and Le, Hieu and Chen, Yihong and Yang, Jiancheng and Wickramasinghe, Udaranga and Fua, Pascal},
  journal={arXiv preprint arXiv:2511.16186},
  year={2025}
}

License

Released under the BSD 3-Clause License, © 2025 EPFL Computer Vision Laboratory (CVLab).

About

repository for PrIntMesh (Precise Intersection Surfaces for 3D Organ Mesh Reconstruction)

Resources

Stars

2 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages