Shape Completion

A modular research platform for 3D shape completion using deep learning. Implements 30+ architectures spanning implicit functions, point cloud completion, diffusion models, transformers, and autoencoders.

Publications

Features

• 30+ Model Architectures: ONet, ConvONet, IF-Net, PCN, SnowflakeNet, ShapeFormer, diffusion models, VAE/VQ-VAE, and more
• Multiple Input Modalities: Point clouds, depth images, RGB-D, simulated Kinect
• Flexible Data Pipeline: Field-based loading with 55+ transforms for augmentation
• Distributed Training: PyTorch Lightning with SLURM/DDP support
• Comprehensive Evaluation: Chamfer, EMD, F1, IoU, FID metrics
• Multi-Backend Rendering: PyRender, Open3D, PyTorch3D, Blender (path tracing)

Quick Start

git clone https://github.com/DLR-RM/shape-completion.git
cd shape-completion
uv sync --inexact --extra ml       # install Python deps (creates .venv automatically)
source .venv/bin/activate          # activate the venv
python libs/libmanager.py install  # compile C++/CUDA libraries (Chamfer, EMD, ...)

# configure dataset paths, then train
# edit conf/dirs/default.yaml first
train -cn shapenet_v1 model.arch=onet

Architecture

shape-completion/
├── models/       # 30+ ML architectures (ONet, ConvONet, diffusion, transformers...)
├── train/        # PyTorch Lightning training infrastructure
├── dataset/      # Data loading with Fields + Transforms system
├── conf/         # Hydra configuration files
├── eval/         # Mesh generation and evaluation metrics
├── inference/    # Real-world inference with preprocessing
├── visualize/    # Mesh extraction (MISE) and multi-backend rendering
├── process/      # Data preprocessing (watertight meshes, Kinect simulation)
├── utils/        # Shared utilities (80+ functions)
└── libs/         # C++/CUDA libraries (Chamfer, EMD, MISE, Kinect sim)

Data Flow

Config (conf/)
    │
    ├── Dataset (dataset/)
    │   ├── Fields → Load data (point clouds, meshes, images)
    │   └── Transforms → Augmentation pipeline
    │
    ├── Model (models/)
    │   └── get_model(cfg) → Architecture selection
    │
    ├── Training (train/)
    │   ├── LitModel → Lightning wrapper
    │   └── Callbacks → Mesh generation, evaluation, visualization
    │
    └── Evaluation (eval/)
        ├── generate → Extract meshes via MISE
        └── mesh_eval → Compute metrics

Available Models

Implicit Shape Completion

Point Cloud Completion

Generative Models

See models/README.md for the complete list of 30+ architectures including DINOv2 variants, point feature extractors, and configuration options.

Installation

Requirements

• uv (recommended) or pip
• CUDA toolkit (for GPU-accelerated libraries)

Python 3.11 is specified in .python-version; uv sync auto-creates the venv with the right interpreter.

1. Clone and install

git clone https://github.com/DLR-RM/shape-completion.git
cd shape-completion
./scripts/bootstrap_dev.sh      # canonical dev/test bootstrap (uses --extra full)
source .venv/bin/activate       # optional, uv run works without manual activation

bootstrap_dev.sh runs uv sync --inexact --extra <extra> and installs Tier-1 custom libs via libs/libmanager.py.
Use ./scripts/bootstrap_dev.sh --extra ml for a lighter profile and --with-tier2-libs when you want PyTorch3D/tiny-cuda-nn/torch-scatter builds too.

2. Choose your extras

Pick one:

uv sync --inexact --extra torch --extra train     # just training
uv sync --inexact --extra ml                      # train + eval + inference
uv sync --inexact --extra full                    # common dependencies
uv sync --inexact --extra all                     # everything

Why --inexact? By default, uv sync removes packages not in the lockfile — including compiled C++/CUDA libs installed via uv pip install (libmanager, compile_cuda_libs.sh). --inexact keeps them.

Individual module extras (dataset, models, eval, inference, visualize, process) are also available for minimal installs.

Note: The diffusion extra installs a custom fork of diffusers with VDM scheduler support. This is required for latent diffusion models.

<details>
<summary>Using pip instead of uv</summary>

python -m venv .venv && source .venv/bin/activate
pip install -e ".[ml]"       # or full, all, etc.

</details>

3. Compile C++/CUDA libraries

Three tiers of compiled dependencies, each building on the previous:

Tier 1: Custom libs (Chamfer, EMD, MISE, Kinect sim, ...)

python libs/libmanager.py install                # all libraries
python libs/libmanager.py install chamfer emd    # or specific ones

install skips already-installed libs. Use upgrade or force-reinstall to overwrite.
See libs/README.md for details and system dependencies (libkinect needs CGAL, Eigen, OpenCV, assimp, libnoise).

Tier 2: PyTorch3D, tiny-cuda-nn, xFormers

These require specific compiler and CUDA arch settings that can't be expressed in the lockfile, so they're installed separately via uv pip install:

# Edit CC/CXX and TORCH_CUDA_ARCH_LIST for your system first
./scripts/compile_cuda_libs.sh

All are optional — the code falls back to pure-PyTorch alternatives when unavailable.

4. Configure paths

Edit conf/dirs/default.yaml with your dataset and log paths:

log: /path/to/training/logs
shapenet_v1_fused: /path/to/shapenet/ShapeNetCore.v1.fused
# ... other datasets

5. Validation lanes

Use the built-in scripts for fast PR checks vs broader nightly checks:

./scripts/test/check_pr.sh       # fast lane: scoped ruff + scoped pyright + CPU-safe pytest subset
./scripts/test/check_all.sh      # full lane: repo-wide ruff/pyright + full pytest --cov (headline informational) + source-only gate
./scripts/test/check_nightly.sh  # broad lane: full pytest with default markers (renderer quarantined)

Quality policy and coverage interpretation live in docs/quality.md. Supported vs experimental workflow status lives in docs/status.md.
GitHub Actions runs the PR lane on pull requests and pushes to main. ./scripts/bootstrap_dev.sh activates the tracked commit-message and pre-push hooks.

Usage

Training

# Basic training
train -cn shapenet_v1

# With model selection
train -cn shapenet_v1 model.arch=conv_onet

# With options
train -cn shapenet_v1 \
    model.arch=conv_onet \
    train.lr=1e-4 \
    train.epochs=200 \
    log.wandb=true

See train/README.md for training options, callbacks, and distributed training.

Evaluation

# Generate meshes from trained model
generate -cn shapenet_v1 model.weights=path/to/model_best.pt

# Evaluate generated meshes
mesh_eval -cn shapenet_v1 test.dir=path/to/meshes

# Combined generation + evaluation
gen_eval -cn shapenet_v1 model.weights=path/to/model_best.pt

See eval/README.md for metrics and evaluation options.

Inference

# On point cloud file
inference input.ply onet -w model.pt -o results/

# On depth image
inference depth.png conv_onet -w model.pt --intrinsic 582.7 582.7 320.8 245.3

# With plane removal (real-world scenes)
inference scene.ply onet -w model.pt --remove_plane --show

See inference/README.md for preprocessing options.

Visualization

# Visualize data loading and augmentation
vis_data -cn shapenet_v1 vis.show=true vis.mesh=true

# Render meshes
render path/to/meshes -o renders/ --method cycles

See visualize/README.md for Generator and Renderer options.

Data Processing

# Generate watertight meshes
make_watertight path/to/ShapeNetCore.v1

# Render Kinect-style depth
render_kinect mesh.obj --out_dir path/to/depth --n_views 1

# Find uncertain regions
find_uncertain_regions -cn shapenet_uncertain

See process/README.md for preprocessing scripts.

Reproducing Paper Results

The checked commands live in docs/reproduction.md and can be rendered with:

python3 scripts/reproduce_experiment.py --list
python3 scripts/reproduce_experiment.py --recipe 3dv2026:ar-uncond-eval \
    --set weights=cvpr_2025/pcd_larm_v1/model_ema.pt \
    --set vae_weights=cvpr_2025_vae/pcd_vqvae_16k_long/model_best.pt

Without --execute, the script prints commands only. This is the recommended public sanity path: it checks the repo-facing command surface without starting long training or evaluation jobs.

C1: Shape Completion with Prediction of Uncertain Regions (IROS 2023)

Project page / arXiv / IEEE

Configs:

Key source files:

Dataset: ShapeNet mugs category (03797390), available on Zenodo. Extract shapenet.tar.gz into a single tree (e.g., <root>/shapenet/03797390/<obj_id>/{samples,blenderproc,mesh,...}).

C1-specific setup (in addition to Setup above):

1. torch_scatter (Tier-2) — required by ConvONet's GridEncoder. Use the prebuilt wheel matching the pinned torch 2.6.0+cu124:

   uv pip install torch-scatter --find-links https://data.pyg.org/whl/torch-2.6.0+cu124.html

   (Skip the full scripts/compile_cuda_libs.sh unless you also want PyTorch3D / Detectron2 / tiny-cuda-nn.)

2. Configure conf/dirs/default.yaml:

   log: /path/to/training-logs
   shapenet_v1_onet: /path/to/zenodo-extract/shapenet

3. taxonomy.json (optional). The Zenodo extract doesn't ship taxonomy.json (ShapeNet license-restricted). Without it, the loader synthesizes a minimal taxonomy — training works, but category labels appear as synset IDs in logs. To restore human-readable names, copy taxonomy.json from upstream ShapeNetCore.v1 into <extract>/shapenet/taxonomy.json.

4. PyTorch3D (optional, paper-exact reproduction only). The shipped mugs.yaml pins inputs.num_points: 2048 (fixed-size subsample) so the config trains without PyTorch3D. The published paper used inputs.num_points: null with PyTorch3D-backed heterogeneous batching. To restore that:

   ./scripts/compile_cuda_libs.sh                # source-build PyTorch3D
   train -cn mugs_paper inputs.num_points=null   # restore variable-size inputs

Commands:

# Preprocess: compute uncertain region labels.
# Skip when using the Zenodo extract (it ships pre-computed *_uncertain.npy labels).
# Only needed if regenerating from raw depth or for non-Zenodo data.
find_uncertain_regions -cn shapenet_uncertain

# Train main model
train -cn mugs_paper

# Train baselines
train -cn mugs_dropout    # MCDropout
train -cn mugs_pssnet     # PSSNet
train -cn mugs_realnvp    # RealNVP

# Evaluate
gen_eval -cn mugs_paper model.weights=path/to/model_best.pt

C2: Combining Shape Completion and Grasp Prediction (Humanoids 2023)

Project page / arXiv / IEEE

Configs:

Key source files:

Dataset: Automatica objects (BOP-style graspable objects), kinect-simulated depth.

Commands:

# Simulate Kinect depth
render_kinect mesh.obj --out_dir path/to/depth --n_views 1

# Train primary model
train -cn humanoids_2023

# Train comparisons
train -cn automatica_2023_kinect_gpt2    # PointTransformer
train -cn shapenet_v1_depth_shapeformer  # ShapeFormer

# Evaluate
gen_eval -cn humanoids_2023 model.weights=path/to/model_best.pt

C3: Evaluating Latent Generative Paradigms (3DV 2026)

Project page / arXiv

Note: C3 configs are named cvpr_2025*, preserving the development naming despite the 3DV 2026 venue.

Configs:

Key source files:

Dataset: ShapeNet v1 (all categories), automatica objects for real-robot evaluation.

Two-stage training:

# Stage 1: Train VAE
train -cn cvpr_2025_vae

# Stage 2: Train LDM (requires VAE checkpoint from stage 1)
# +vae_weights is relative to dirs.log (set in conf/dirs/default.yaml)
train -cn cvpr_2025 +vae_weights=cvpr_2025_vae/<run_name>/model_best.pt

# Evaluate (must pass +vae_weights again — LDM checkpoints don't include VAE weights)
gen_eval -cn cvpr_2025 model.weights=path/to/ldm/model_best.pt \
    +vae_weights=cvpr_2025_vae/<run_name>/model_best.pt

For the VQ-VAE/AR recipes and the normalization overrides used by the released runs, prefer the checked recipes in docs/reproduction.md.

C4: Instance-aware 3D completion

This codebase also contains the instance-completion functionality used for joint instance segmentation and shape completion on tabletop scenes. The public entry point is conf/tabletop_inst_seg_pcd_3d.yaml, with model implementations in models/src/dinov2.py and dataset support in dataset/src/tabletop.py.

The checked recipes cover:

python3 scripts/reproduce_experiment.py --list --suite inst3d2026

This includes joint training, pipeline training, pile fine-tuning/evaluation, and BOP scene evaluation. Pile recipes use inputs.type=kinect_sim, data.train_ds=[tabletop_pile], and test.no_aug=True. BOP recipes use data.frame=cam; use camera-frame checkpoints for those runs.

Configuration

All experiments use Hydra configs from conf/. The base config (config.yaml) defines all options with defaults.

# Use a config
train -cn shapenet_v1

# Override values
train -cn shapenet_v1 model.arch=conv_onet train.lr=1e-4

# Multi-run sweep
train -cn shapenet_v1 model.arch=onet,conv_onet,if_net --multirun

See conf/README.md for the complete configuration schema.

Custom Libraries

Install via python libs/libmanager.py install. See libs/README.md for per-library build dependencies.

Modules

Example Workflows

Train ONet on ShapeNet

# 1. Train
train -cn shapenet_v1 model.arch=onet train.epochs=100

# 2. Generate meshes
generate -cn shapenet_v1 \
    model.weights=logs/shapenet_v1/onet/version_0/model_best.pt

# 3. Evaluate
mesh_eval -cn shapenet_v1 \
    test.dir=logs/shapenet_v1/onet/version_0/meshes

Real-World Inference

# Inference on Kinect scan with plane removal
inference kinect_scene.ply conv_onet \
    -w model.pt \
    --remove_plane \
    --cluster \
    -o completed/

Prepare Custom Dataset

# 1. Generate watertight meshes
make_watertight path/to/meshes --out_dir path/to/watertight

# 2. Create config
# conf/my_dataset.yaml with data paths

# 3. Train
train -cn my_dataset model.arch=onet

FAQ

How to install CUDA without root?

wget https://developer.download.nvidia.com/compute/cuda/12.4.0/local_installers/cuda_12.4.0_550.54.14_linux.run
chmod +x cuda_12.4.0_550.54.14_linux.run
./cuda_12.4.0_550.54.14_linux.run --toolkit --toolkitpath=/path/to/custom/directory

How to use multiple GPUs?

# Local multi-GPU
CUDA_VISIBLE_DEVICES=0,1,2,3 train -cn shapenet_v1

# SLURM
srun --gres=gpu:4 train -cn shapenet_v1

How to resume training?

train -cn shapenet_v1 train.resume=true
# or specify checkpoint
train -cn shapenet_v1 model.checkpoint=path/to/last.ckpt

How to use a different attention backend?

train -cn shapenet_v1 model.attn_backend=xformers  # or torch, einops

Citation

If you use this codebase, please cite the relevant paper(s):

@inproceedings{humt2023shape,
  author    = {Humt, Matthias and Winkelbauer, Dominik and Hillenbrand, Ulrich},
  title     = {Shape Completion with Prediction of Uncertain Regions},
  booktitle = {2023 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)},
  year      = {2023},
  pages     = {1215--1221},
  doi       = {10.1109/IROS55552.2023.10342487},
}

@inproceedings{humt2023combining,
  author       = {Humt, Matthias and Winkelbauer, Dominik and Hillenbrand, Ulrich and B{\"a}uml, Berthold},
  title        = {Combining Shape Completion and Grasp Prediction for Fast and Versatile Grasping with a Multi-Fingered Hand},
  booktitle    = {2023 IEEE-RAS 22nd International Conference on Humanoid Robots (Humanoids)},
  year         = {2023},
  pages        = {1--8},
  doi          = {10.1109/Humanoids57100.2023.10375210},
  organization = {IEEE},
}

@inproceedings{humt2026generative,
  author        = {Humt, Matthias and Hillenbrand, Ulrich and Triebel, Rudolph},
  title         = {Evaluating Latent Generative Paradigms for High-Fidelity {3D} Shape Completion from a Single Depth Image},
  booktitle     = {International Conference on 3D Vision (3DV)},
  year          = {2026},
  doi           = {10.48550/arXiv.2511.11074},
  eprint        = {2511.11074},
  archiveprefix = {arXiv},
  primaryclass  = {cs.CV},
}

License

This project is licensed under the MIT License. See LICENSE for details.