README.md

metadata

license: mit
library_name: cod-vae
pipeline_tag: feature-extraction
tags:
  - 3d
  - shape-reconstruction
  - autoencoder
  - vae
  - occupancy

COD-VAE 16 x 8 (small)

A compact COD-VAE that compresses a 3D shape into 16 latent vectors of 8 dimensions = 128 numbers, and decodes them back into an occupancy field. Same latent shape as cod-vae-16x8, but a ~4x smaller network trained for fast decoding: ~35M parameters instead of 188M, with a 14.2M-parameter decode path instead of 90M.

Note: although the latent shape matches cod-vae-16x8, the two models define different latent spaces — latents from one cannot be decoded with the other.

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.

Architecture vs cod-vae-16x8

	cod-vae-16x8	this model
embed dim / heads	512 / 8	256 / 4
encoder	4 blocks x 3 layers	3 blocks x 3 layers
latent decoder layers	12	6
refinement decoder layers	12	8
total parameters	188M	~35M
decode-path parameters	90M	14.2M

Decode speed (H100, float32 + TF32)

	cod-vae-16x8	this model
batch-1 latency (latents -> triplanes)	4.11 ms	2.69 ms
batch-32 throughput	2,720 shapes/s	5,358 shapes/s
forward+backward, batch 32 x 2048 queries	53.0 ms	25.4 ms

The dense 128^3 query pass (~3.7 ms) is unchanged; it depends only on the triplane query head, not on model width or depth.

Usage

import trimesh
from cod_vae import CODVAE

vae = CODVAE.from_pretrained("TimSchneider42/cod-vae-16x8-small")

mesh = trimesh.load("bunny.obj", force="mesh")
latent, transform = vae.encode_mesh(mesh, return_transform=True)   # (16, 8)
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 schedule (100 + 100 epochs); only the batch layout differs from cod-vae-16x8 because the smaller model needs fewer GPUs:

	stage 1 (autoencoder)	stage 2 (latent VAE)
epochs	100	100
batch	64 per GPU x 4 GPUs = 256	256 per GPU x 2 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.8489	0.8085
MNIST3D (embossed digits)	128	0.9105	0.8698

For reference, the full-size cod-vae-16x8 reaches 0.8733 / 0.8347 on ABC and 0.9231 / 0.8829 on MNIST3D — the size and speed here cost about 0.01–0.03 IoU. 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}
}