TimSchneider42/cod-vae-8x32

COD-VAE 8 x 32

A COD-VAE that compresses a 3D shape into 8 latent vectors of 32 dimensions = 256 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-8x32")

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

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

source	shapes	query pools per shape
ShapeNet (3DShape2VecSet, 55 synsets)	48,597	500k volume + 500k near-surface
tactile-mnist-abc-dataset-small	50,000	500k + 500k
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 recipe.

Training recipe

Both stages follow the reference implementation; see 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.8868	0.8513
MNIST3D (embossed digits)	128	0.9304	0.8930

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