README.md

---
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}
}
```