Self-supervised Learning of Fine-to-Coarse Cuboid Shape Abstraction

Gregor Kobsik, Morten Henkel, Yanjiang He, Victor Czech, Tim Elsner, Isaak Lim, and Leif Kobbelt
Visual Computing Institute, RWTH Aachen University, Germany

Eurographics 2026 — May 4–8, Aachen, Germany

Overview

This repository contains the official implementation of Self-supervised Learning of Fine-to-Coarse Cuboid Shape Abstraction. The method learns to abstract 3D shapes (point clouds) into a compact set of oriented cuboids in a fully self-supervised manner, progressively reducing the number of primitives from a large initial set to a compact target abstraction during training.

The model uses a transformer-based encoder built on Irregular Latent Grids (ILG) to encode input point clouds into latent features, and a transformer-based cuboid decoder that predicts rotation, translation, scale, and existence probability for each cuboid. An abstraction loss gradually drives the active cuboid count from start_cuboids down to target_cuboids over the course of training.

Method

The model operates in two stages:

1. Encoding: A point cloud (2048 surface points) is encoded by Encoder3DILG, which uses farthest-point sampling, local PointConv aggregation, and a Vision Transformer to produce N latent feature vectors at irregular 3D positions.

2. Decoding: CuboidDecoder attends over the latent features and predicts per-cuboid parameters — rotation (3×3), translation (3), scale (3), and existence probability (scalar) — for up to start_cuboids cuboids simultaneously.

Losses:

• Surface reconstruction loss: bidirectional distance between cuboid surface samples and input surface points, weighted by cuboid area.
• Volume reconstruction loss: bidirectional distance between cuboid volume samples and input volume points, weighted by cuboid volume.
• Abstraction loss: a schedule-based cross-entropy loss that drives the expected number of active cuboids from start_cuboids → target_cuboids using a half-cosine annealing schedule.

Active cuboids are pruned each validation epoch based on their empirical existence probability across the dataset.

---

Installation

pip install torch lightning pytorch3d torch-cluster einops fpsample pl_bolts

---

Data

Point Clouds

We use the presampled ShapeNet point clouds provided by:

Kaizhi Yang and Xuejin Chen. Unsupervised Learning for Cuboid Shape Abstraction via Joint Segmentation from Point Clouds. SIGGRAPH 2021.

Download the data from their repository: https://github.com/SilenKZYoung/CuboidAbstractionViaSeg

Supported categories: airplane, chair, table. The human category uses a separate dataset (see below).

Human Category

The human category is sampled and precomputed using a subset of the DFAUST dataset:

Federica Bogo, Javier Romero, Gerard Pons-Moll, and Michael J. Black. Dynamic FAUST: Registering Human Bodies in Motion. CVPR 2017.

Dataset: https://dfaust.is.tue.mpg.de/

Preprocessing

Use data/pre_processing.py to prepare the inputs expected by the dataloader from the raw meshes. The script computes two sets of samples per shape using ShapeNetCore v2 (shape categories) and DFAUST (human) as ground truth:

Split	Sampling strategy	Count
Volume	Uniformly distributed in the unit cube	50 000
Near-surface	Sampled on the mesh surface, perturbed by ±5%	50 000

---

Training

python main.py \
    --category airplane \
    --start_cuboids 128 \
    --target_cuboids 7 \
    --num_epochs 1000 \
    --batch_size 16

Key arguments:

Argument	Default	Description
--category	airplane	ShapeNet category
--start_cuboids	128	Initial number of cuboids
--target_cuboids	7	Target number of cuboids after abstraction
--num_epochs	1000	Training epochs
--batch_size	16	Batch size
--embedding_size	128	Latent feature dimension
--num_latents	128	Number of ILG latent points
--num_layers	6	Transformer depth
--num_heads	4	Attention heads
--loss_w_surface	1e0	Surface loss weight
--loss_w_volume	1e1	Volume loss weight
--loss_w_abstract	1e-3	Abstraction loss weight

Training logs and checkpoints are saved under logs/ via TensorBoard.

---

Pre-trained Checkpoints

Pre-trained checkpoints for all categories are available as pre_trained.zip in the Releases section of this repository.

---

Evaluation

Metrics reported at test time:

• Surface-CD: Chamfer distance between the cuboid abstraction and the input surface.
• Volume-IoU: Intersection-over-Union between the cuboid abstraction and the voxelized shape.
• Cuboids-Num: Mean number of active cuboids per shape.

---

Visualization

Use visualization.ipynb to load a trained checkpoint and render cuboid abstractions for individual shapes.

---

Citation

If you use this code, please cite our paper:

@article{kobsik2026cuboid,
  title={Self-supervised Learning of Fine-to-Coarse Cuboid Shape Abstraction},
  author={Kobsik, Gregor and Henkel, Morten and He, Yanjiang and Czech, Victor and Elsner, Tim and Lim, Isaak and Kobbelt, Leif},
  year={2026},
  journal={Computer Graphics Forum},
  volume={45},
  number={2},
}

---

License

This project is released for research and non-commercial use.