cod-vae-32x4

COD-VAE 32 x 4

A COD-VAE that compresses a 3D shape into 32 latent vectors of 4 dimensions = 128 numbers, and decodes them back into an occupancy field.

Trained with 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

import trimesh
from cod_vae import CODVAE

vae = CODVAE.from_pretrained("TimSchneider42/cod-vae-32x4")

mesh = trimesh.load("bunny.obj", force="mesh")
latent, transform = vae.encode_mesh(mesh, return_transform=True)   # (32, 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:

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:

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

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 recipe.

Training recipe

Both stages follow the reference implementation; see TRAINING.md for the full guide and the exact commands.

Held-out reconstruction quality

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:

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