| --- |
| license: mit |
| library_name: cod-vae |
| pipeline_tag: feature-extraction |
| tags: |
| - 3d |
| - shape-reconstruction |
| - autoencoder |
| - vae |
| - occupancy |
| --- |
| |
| # COD-VAE 8 x 4 |
|
|
| A [COD-VAE](https://arxiv.org/abs/2503.08737) that compresses a 3D shape into |
| **8 latent vectors of 4 dimensions = 32 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. |
|
|
| 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-8x4") |
| |
| mesh = trimesh.load("bunny.obj", force="mesh") |
| latent, transform = vae.encode_mesh(mesh, return_transform=True) # (8, 4) |
| 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.7269 | 0.7479 | |
| | MNIST3D (embossed digits) | 128 | 0.7271 | 0.7775 | |
|
|
| 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} |
| } |
| ``` |
|
|
