Sparse 3D Shape Completion via SDF Diffusion

This repository provides an EDM diffusion-based pipeline for 3D shape completion operating natively on sparse Signed Distance Fields (SDFs) via spconv.

Overview

• End-to-End Pipeline: Reproducible scripts for data preprocessing, training, sampling, and evaluation.
• Adjustable Sparse TSDFs: An intuitive extension of Truncated Signed Distance Fields (TSDF) that allows for dynamic voxel distance cutoffs via the --band_width parameter, converted into sparse spconv tensors.
• Dense Baseline Comparison: Includes an implementation of the dense DiffComplete architecture, explicitly adapted to the EDM diffusion formulation to enable direct comparison.

Installation

Ensure you have a CUDA-capable environment (tested on CUDA 12.6) and Python 3.10+.

# Clone the repository
git clone https://github.com/defy01/sparse-3d-completion.git
cd sparse-3d-completion

# Optionally create and activate a virtual environment
python3 -m venv venv
source venv/bin/activate

# Install dependencies
pip install -r requirements.txt

Note: If using a different CUDA version, adjust the spconv-cu126 and torch requirement in requirements.txt accordingly.

Pre-trained Models

Trained model weights for the bidirectional sparse architecture and the dense baseline can be downloaded here:

• Download Checkpoints (Google Drive)

Place the downloaded .pt files into the checkpoints/ directory maintaining the structure expected by the configs (e.g., checkpoints/bidir-xs_64/ckpt_0100000.pt).

Execution Examples

Configuration files for training and evaluation are located in the configs/ directory. You can override any YAML parameter directly from the command line using the --override flag.

1. Data Preprocessing

Download and process a curated subset of Objaverse furniture meshes (including chairs, tables, sofas, cabinets, beds, and benches) into normalized SDF grids.

python scripts/preprocessing/objaverse_to_sdf.py \
    --out_dir data/sdf64 \
    --resolution 64 \
    --conversion_processes 4 \
    --splits_out data/splits_labels.json

To additionally precompute the holed SDFs for the validation (or test) set, run:

python scripts/preprocessing/generate_eval_data.py \
    --src_dir data/sdf64 \
    --split_json data/splits_labels.json \
    --out_dir data/val64 \
    --split val \
    --band_width 4

Note: To generate the testing set, change --split val to --split test and adjust the --out_dir accordingly.

2. Training

Launch distributed mixed-precision training using Hugging Face accelerate.

accelerate launch train.py \
    --config configs/train/bidir-xs_64.yaml \
    --override batch_size=10

(Hardware Note: Training the network at a resolution of 64 with a batch size of 10 requires a GPU with at least 24 GB of VRAM, e.g., RTX 3090 / A5500).

3. Sampling

Load a complete SDF shape, simulate random missing regions (holes), and generate the completed 3D shape using a trained checkpoint.

python scripts/sample.py \
    --ckpt checkpoints/bidir-xs_64/ckpt_0100000.pt \
    --input_sdf data/sdf64/example_shape.npy \
    --out_dir results/ \
    --band_width 4 \
    --guidance_scale 1.25

(Hardware Note: Inference has a significantly lower memory footprint than training; any standard GPU with 8 GB of VRAM should be sufficient for sampling).

4. Evaluation

Evaluate a trained checkpoint to compute Global/Masked IoU, Chamfer Distance, and F-Score.

python scripts/evaluate.py \
    --config configs/eval/bidir-xs_64.yaml \
    --override ckpt=checkpoints/bidir-xs_64/ckpt_0100000.pt

Acknowledgements

This project is built upon and inspired by several foundational works. We sincerely thank the authors for their research and open-source contributions:

• EDM: Elucidating the Design Space of Diffusion Models (Karras et al., 2022)
• DiffComplete: DiffComplete: Diffusion-based Generative 3D Shape Completion (Chu et al., 2023)
• ControlNet-XS: ControlNet-XS: Rethinking the Control of Text-to-Image Diffusion Models as Feedback-Control Systems (Zavadski et al., 2024)