Xiang Xu, Joseph Lambourne, Pradeep Jayaraman, Zhengqing Wang, Karl Willis, and Yasutaka Furukawa
We present a diffusion-based generative approach that directly outputs a CAD B-rep. BrepGen uses a novel structured latent geometry to encode the CAD geometry and topology. A top-down generation approach is used to denoise the faces, edges, and vertices.
- Linux
- Python 3.9
- CUDA 11.8
- PyTorch 2.2
- Diffusers 0.27
Install PyTorch and other dependencies:
conda create --name brepgen_env python=3.9 -y
conda activate brepgen_env
pip install -r requirements.txt
pip install chamferdist
If chamferdist
fails to install here are a few options to try:
-
If there is a CUDA version mismatch error, then try setting the
CUDA_HOME
environment variable to point to CUDA installation folder. The CUDA version of this folder must match with PyTorch's version i.e. 11.8. -
Try building from source.
Install OCCWL following the instruction here. If conda is stuck in "Solving environment..." there are two options to try:
-
Try using
mamba
as suggested in occwl's README. -
Install pythonOCC: https://github.com/tpaviot/pythonocc-core?tab=readme-ov-file#install-with-conda and occwl manually:
pip install git+https://github.com/AutodeskAILab/occwl
.
Download ABC STEP files (100 folders).
Download Furniture Data. JSON file contains object UIDs that refer to the original STEP files.
The faces, edges, and vertices need to be extracted from the STEP files.
Process the B-reps (under data_process
folder):
sh process.sh
Remove repeated CAD models (under data_process
folder, default is 6 bit
):
sh deduplicate.sh
You can download the deduplicated files for DeepCAD, and ABC.
Train the surface and edge VAE (wandb for logging):
sh train_vae.sh
Train the latent diffusion model (change path to previously trained VAEs):
sh train_ldm.sh
--cf
classifier-free training for the Furniture dataset.
--data_aug
randomly rotate the CAD model during training (optional).
Randomly generate B-reps from Gaussian noise, both STEP and STL files will be saved:
python sample.py --mode abc
This will load the settings in eval_config.yaml
. Make sure to update model paths to the correct folder.
Run this script for evaluation (change the path to generated data folder, with at least 3,000 samples):
sh eval.sh
This computes the JSD, MMD, and COV scores. Please also download sampled point clouds for test set.
We also provide the individual checkpoints trained on different datasets.
Source Dataset | ||
---|---|---|
DeepCAD | vae model | latent diffusion model |
ABC | vae model | latent diffusion model |
This research is partially supported by NSERC Discovery Grants with Accelerator Supplements and DND/NSERC Discovery Grant Supplement, NSERC Alliance Grants, and John R. Evans Leaders Fund (JELF). We also thank Onshape for their support and access of the publicly available CAD models.
If you find our work useful in your research, please cite the following paper
@article{xu2024brepgen,
title={BrepGen: A B-rep Generative Diffusion Model with Structured Latent Geometry},
author={Xu, Xiang and Lambourne, Joseph G and Jayaraman, Pradeep Kumar and Wang, Zhengqing and Willis, Karl DD and Furukawa, Yasutaka},
journal={arXiv preprint arXiv:2401.15563},
year={2024}
}