A generalizable flow matching-based 3D reassembly method trained on 1.9 Million fractures, enabling precise real-world fragment pose alignment. 😊Achieves strong performance across extensive benchmarks, concise code with efficient performance.
Sihang Li* · Zeyu Jiang* · Grace Chen† · Chenyang Xu† · Siqi Tan · Xue Wang · Irving Fang · Kristof Zyskowski · Shannon McPherron · Radu Iovita · Chen Feng✉ · Jing Zhang✉
*, † Equal contribution ✉ Corresponding author
2025/03/25: We release the GARF, which achieves state-of-the-art performance across a diverse range of synthetic and real-world benchmarks. Try our demo on your own data!
We recommend using uv to manage the dependencies. Follow the instructions here to install uv. Then, simply run
uv sync
uv sync --extra post
source ./venv/bin/activateto install the dependencies and activate the virtual environment. Please be noted that flash-attn requires CUDA 12.0 or above and pytorch3d may need GPU available when installing it, or it will compile with no GPU support.
If you encounter any issue, you may try to re-install after removing the virtual environment at .venv and doing uv clean to remove the cache.
We will soon provide the script to process the raw Breaking Bad dataset into our hdf5 format, right now, you can directly download our processed dataset from following links. Fractura dataset will be released soon.
| Dataset | Link |
|---|---|
| Breaking Bad | OneDrive Google Drive |
| Breaking Bad Volume Constrained | OneDrive Google Drive |
| Breaking Bad Other | OneDrive Google Drive |
We provide the evaluation script in scripts/eval.sh, which looks like this:
EXPERIMENT_NAME="everyday_vol_one_step_init"
DATA_ROOT="../breaking_bad_vol.hdf5"
DATA_CATEGORIES="['everyday']"
CHECKPOINT_PATH="output/GARF.ckpt"
HYDRA_FULL_ERROR=1 python eval.py \
seed=42 \
experiment=denoiser_flow_matching \
experiment_name=$EVAL_NAME \
loggers=csv \
loggers.csv.save_dir=logs/GARF \
trainer.num_nodes=1 \
trainer.devices=[0] \
data.data_root=$DATA_ROOT \
data.categories=$DATA_CATEGORIES \
data.batch_size=64 \
ckpt_path=$CHECKPOINT_PATH \
++data.random_anchor=false \
++model.inference_config.one_step_init=trueIf you want to evaluation on our diffusion variant, you can simply change the experiment to denoiser_diffusion.
After running the script, the evaluation results will be stored in logs/GARF/everyday_vol_one_step_init/. In this folder, you could refer to metrics.csv for the evaluation results, and there's a json_results folder which contains the reassembly results for all test samples.
Currently, if you run the evaluation using multi-gpu, the json_results maybe incomplete, so we recommend using single gpu for evaluation. We will fix this issue in the future.
The training process is quite similar to the evaluation process. While you could directly run the training script provided below, we recommend getting familiar with our project structure and config system first.
NUM_NODES=4
DATA_ROOT="../breaking_bad_vol.hdf5"
DATA_CATEGORIES="['everyday']"
CHECKPOINT_PATH="output/pretraining.ckpt"
python train.py \
experiment=pretraining_frac_seg \
data.categories=$DATA_CATEGORIES \
trainer.num_nodes=$NUM_NODES \
data.data_root=$DATA_ROOT \
ckpt_path=$CHECKPOINT_PATH # to resume trainingThe difference here is that we will use the pretrained feature extractor to initialize the model, and we have to change the experiment into our flow-matching training.
NUM_NODES=4
DATA_ROOT="../breaking_bad_vol.hdf5"
DATA_CATEGORIES="['everyday']"
FEATURE_EXTRACTOR_CKPT="output/pretraining.ckpt"
CHECKPOINT_PATH="output/GARF.ckpt"
python train.py \
experiment=denoiser_flow_matching \
data.categories=$DATA_CATEGORIES \
trainer.num_nodes=$NUM_NODES \
data.data_root=$DATA_ROOT \
model.feature_extractor_ckpt=$FEATURE_EXTRACTOR_CKPT \
ckpt_path=$CHECKPOINT_PATH # to resume trainingTo start fine-tuning, the very first thing you need to do is to prepare your own dataset. The dataset should be in the same format as the Breaking Bad dataset, and you can use our provided script to convert it into hdf5 format. Then, you can run the following example script to start fine-tuning.
python train.py \
experiment=finetune \
experiment_name=finetune \
data.categories="['egg']" \
project_name="GARF" \
trainer.num_nodes=$NUM_NODES \
data.data_root=./finetune_egg.hdf5 \
data.num_workers=8 \
data.batch_size=32 \
data.multi_ref=True \
tags='["finetune", 'egg']' \
ckpt_path=./xxx \
finetuning=true.
├── assembly
│ ├── backbones # Backbones used for feature extraction
│ ├── data # Data processing module
│ └── models
│ ├── denoiser # Denoising models
│ └── pretraining # Pretraining module
├── configs # Configuration files directory
├── eval.py # Evaluation script
├── train.py # Training script
└── vis.py # Visualization scriptAll the configuration files are stored in the configs folder. The config system is based on Hydra, which allows you to easily modify the configurations by changing the YAML files. You can also override the configurations by passing command line arguments when running the script. We hugely utilize the config system for the initialization of all the modules. You could refer to configs/models to see the configuration files for different models. The configs/experiments folder serves as a global override configuration for the training and evaluation scripts.
After running the evaluation, per sample transformations will be saved in logs/GARF/{EXPERIMENT_NAME}/json_results/. Using the transformation saved in the json, you can firstly apply the inverse of gt transformation to the fragments to get the model input, and then apply the model output transformations to the fragments to get the the final output. We'll soon provide a script and sample to visualize the results.
| Model Name | Model | Note |
|---|---|---|
| GARF-mini | GARF-mini-E-FM | pretrained on everyday subset of Breaking Bad with Flow-matching model. |
| GARF-mini-diffusion | GARF-mini-E-Diff | replace the Flow-matching model with Diffusion model |
| GARF | GARF-EAO-FM | large-scale trained on everyday+artifact+other subsets of Breaking Bad for both backbone and Flow-matching (cost most time!) |
| Dataset | Subset | Model | RMSE(R) ↓ | RMSE(T) ↓ | Part Accuracy ↑ |
|---|---|---|---|---|---|
| Breaking Bad Vol | Everyday | GARF | 5.32 | 1.14 | 95.68% |
| Breaking Bad Vol | Everyday | GARF-mini | 6.68 | 1.34 | 94.77% |
| Breaking Bad Vol | Artifact | GARF | 3.64 | 0.88 | 96.78% |
| Breaking Bad Vol | Artifact | GARF-mini | 7.67 | 1.77 | 93.34% |
For frequently asked questions, please refer to our GitHub Issues page. You can search existing issues or create a new one if you don't find an answer to your question.
If you find this project useful, please consider citing our paper:
@article{Li2025GARF,
title={GARF: Learning Generalizable 3D Reassembly for Real-World Fractures},
author={Sihang Li and Zeyu Jiang and Grace Chen and Chenyang Xu and Siqi Tan and Xue Wang and Irving Fang and Kristof Zyskowski and Shannon McPherron and Radu Iovita and Chen Feng and Jing Zhang},
year={2025},
journal={arXiv preprint arXiv:2504.05400},
}Our codebase and method implementation are based on the excellent work by PuzzleFusion++ and PointTransformerV3, which provided valuable foundations and insights.
This project is licensed under the GPL License. See the LICENSE file for details.
We gratefully acknowledge the Physical Anthropology Unit, Universidad Complutense de Madrid for providing access to the human skeletons under their curation. This work was supported in part through NSF grants 2152565, 2238968, 2322242, and 2426993, and the NYU IT High Performance Computing resources, services, and staff expertise.
@article{li2025garf,
title={GARF: Learning Generalizable 3D Reassembly for Real-World Fractures},
author={Li, Sihang and Jiang, Zeyu and Chen, Grace and Xu, Chenyang and Tan, Siqi and Wang, Xue and Fang, Irving and Zyskowski, Kristof and McPherron, Shannon P and Iovita, Radu and Feng, Chen and Zhang, Jing},
year={2025},
journal={arXiv preprint arXiv:2504.05400},
}
