File size: 4,293 Bytes
2ce06a8 74a17ce 2ce06a8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | ---
license: mit
library_name: cod-vae
pipeline_tag: feature-extraction
tags:
- 3d
- shape-reconstruction
- autoencoder
- vae
- occupancy
---
# COD-VAE 64 x 32
A [COD-VAE](https://arxiv.org/abs/2503.08737) that compresses a 3D shape into
**64 latent vectors of 32 dimensions = 2048 numbers**, and decodes them back
into an occupancy field.
Trained with [`cod-vae`](https://github.com/TimSchneider42/cod-vae), a PyTorch/JAX reimplementation of COD-VAE
(Cho et al., ICCV 2025). The weights are a self-contained npz and load with either
backend.
This repository is the **default model** of the `cod-vae` package: it holds the same
weights as [`TimSchneider42/cod-vae-64x32`](https://huggingface.co/TimSchneider42/cod-vae-64x32),
the largest configuration of the grid, under the short name `TimSchneider42/cod-vae`. The Hub
has no aliasing between repositories, so the two are independent copies of one model.
Stage 1 ran the full 100 epochs and stage 2 another 100, following the
reference schedule throughout.
## Usage
```python
import trimesh
from cod_vae import CODVAE
vae = CODVAE.from_pretrained("TimSchneider42/cod-vae")
mesh = trimesh.load("bunny.obj", force="mesh")
latent, transform = vae.encode_mesh(mesh, return_transform=True) # (64, 32)
reconstruction = vae.decode_mesh(latent, transform=transform) # trimesh.Trimesh
```
Latents can also be computed from raw surface point clouds and decoded at arbitrary
query points:
```python
latents = vae.encode(points) # (N, 3) in [-1, 1]^3
logits = vae.decode(latents, queries) # occupancy logits, positive inside
volume = vae.decode_volume(latents, resolution=128) # dense logit grid
```
Install with `pip install cod-vae[torch,hub]` (or `cod-vae[jax,hub]`).
## Training data
A merged dataset of 110,077 shapes, built with the `cod-vae-dataset` tool:
```bash
cod-vae-dataset data/merged --vecset path/to/shapenet_vecset_root
cod-vae-dataset data/merged \
--hf abc=TimSchneider42/tactile-mnist-abc-dataset-small:0.24435897 --hf-split train \
--num-vol 500000 --num-surface 250000
cod-vae-dataset data/merged \
--hf mnist3d=TimSchneider42/tactile-mnist-mnist3d --hf-split train \
--num-vol 50000 --num-surface 25000
```
| source | shapes | query pools per shape |
|---|---|---|
| ShapeNet (3DShape2VecSet, 55 synsets) | 48,597 | 500k volume + 500k near-surface |
| [tactile-mnist-abc-dataset-small](https://huggingface.co/datasets/TimSchneider42/tactile-mnist-abc-dataset-small) | 50,000 | 500k + 500k |
| [tactile-mnist-mnist3d](https://huggingface.co/datasets/TimSchneider42/tactile-mnist-mnist3d) | 11,480 | 50k + 50k |
Only the training splits are used; the ABC and MNIST3D pool sizes are scaled to the
geometric complexity of each source. Meshes are preprocessed with the original authors'
[sdf_gen](https://github.com/1zb/sdf_gen) recipe.
## Training recipe
Both stages follow the reference implementation; see
[TRAINING.md](https://github.com/TimSchneider42/cod-vae/blob/main/TRAINING.md) for the full guide and the exact commands.
| | stage 1 (autoencoder) | stage 2 (latent VAE) |
|---|---|---|
| epochs | 100 | 100 |
| batch | 32 per GPU x 2 accumulation x 4 GPUs = 256 | 128 per GPU x 4 GPUs = 512 |
| learning rate | 1e-4, scaled by effective batch / 256 | same, halved at epochs 60/70/80/90 |
| dataset repeat | 8 per epoch | 8 per epoch |
| precision | float32 with TF32 matmuls | same |
## Held-out reconstruction quality
| source | held-out shapes | volume IoU | near-surface accuracy |
|---|---|---|---|
| ABC (CAD parts) | 128 | 0.9316 | 0.8932 |
| MNIST3D (embossed digits) | 128 | 0.9531 | 0.9207 |
Measured on the test splits of ABC and MNIST3D, which are disjoint from training.
Volume IoU compares `decode(latents, queries) > 0` against ground-truth occupancy on
uniformly sampled query points; near-surface accuracy uses points sampled around the
surface.
## Citation
The model architecture and training recipe are from:
```bibtex
@inproceedings{cho2025cod,
author={Cho, In and Yoo, Youngbeom and Jeon, Subin and Kim, Seon Joo},
title={Representing 3D Shapes with 64 Latent Vectors for 3D Diffusion Models},
booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
year={2025}
}
```
|