Modelpublic

TimSchneider42/cod-vae-4x8

sourceHugging Facemitupdated 2mo agoView on Hugging Face

0likes

Model Card

COD-VAE 4 x 8

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

python

import trimesh
from cod_vae import CODVAE

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

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

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	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.7434	0.7576
MNIST3D (embossed digits)	128	0.7624	0.7925

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