MeshWalker for 3D Object Registration

Extending MeshWalker — a random-walk-based deep learning model for 3D mesh understanding — from single-mesh classification to pairwise 3D object registration: given two meshes of the same object (up to a linear transformation), predict the transformation matrix that aligns them.

This was a DTU Special Course project by Amir Kfir, built on top of the original MeshWalker codebase by Alon Lahav and Ayellet Tal (SIGGRAPH Asia 2020). The full write-up, including methodology, experiments, and results, is available in report/MeshWalker_Object_Registration_Report.pdf.

Motivation

3D object registration — finding the transformation that aligns two similar 3D objects — is a core task in computer-vision-for-robotics. Classical algorithms are sensitive to hyperparameters and trade off accuracy against runtime; learning-based approaches (e.g. NrtNet, DWC) tend to process full point clouds, which is computationally heavy.

Instead, this project feeds the network sparse 3D trajectories generated by MeshWalker's random-walk sampling over the mesh surface, and trains a sequential model to regress the alignment transform directly from a pair of walks — avoiding full point-cloud processing.

Approach

1. Walk generation — MeshWalker's random-walk algorithm samples a trajectory of vertices from each of the two input meshes.
2. Feature extraction — per-point FC layers lift each walk into a higher-dimensional representation.
3. Sequence modeling — the two high-dimensional trajectories are processed by a sequential network (see architectures below).
4. Regression — a final module predicts the linear transformation matrix aligning the pair.

Sequential architectures (RNN / GRU / LSTM)

Two fusion strategies were compared: early merge, where the two walks are concatenated before the sequential block, and late merge, where each walk is encoded independently before being concatenated for the final regression.

Transformer architecture

An attention-based alternative was also implemented: the first encoder block takes one walk as the query and the other as key/value, followed by further self-attention blocks and an MLP decoder. The intuition is that attention can align corresponding but unordered points between the two walks — similar to aligning differently-ordered sentences in translation.

Results

Multiple architectures, sequence lengths, and augmentation ranges were evaluated (see the report for the full table and discussion).

Model Type	Walk Size	Scaling	Translation	Rotation	Seq. Blocks	MSE
Transformers	100	1–1	±0	±12°	4	0.00619
LSTM-single	1000	1–1	±0	±12°	4	0.00641
LSTM-single	300	1–1	±0	±12°	4	0.0065
LSTM-single	100	1–1	±0	±12°	4	0.0063
LSTM-dual	100	1–1	±0	±12°	4	0.0099
GRU-dual	100	1–1	±0	±12°	4	0.0099
LSTM-single	100	1–1	±0	±60°	4	0.121
Transformers	100	1–1	±0	±60°	4	0.116
GRU-dual	100	0.5–2	±5	±360°	4	2.29

Takeaways: early-merge fusion consistently outperformed late-merge, and walk length had little effect on accuracy. The network converged reasonably well on small rotation ranges but could not learn full-range rotation to a usable degree of accuracy — the proof of concept did not reach a production-ready registration accuracy. The negative result is still informative: since MeshWalker trajectories are known to be sufficient for single-mesh classification, the bottleneck is more likely in the sequence-pair fusion architecture than in the trajectory representation itself. Promising directions for follow-up are outlined in the report (non-sequential encoders, covariance-based trajectory features, residual sequential blocks).

Repository structure

This repo is a fork of the original MeshWalker implementation. Files added for the object-registration extension:

File	Purpose
train_object_registration.py	Training entry point for the registration task
evaluate_object_registration.py	Evaluation / accuracy reporting for a trained registration model
rnn_object_registration_model.py	Dual-input RNN/GRU/LSTM registration network (early & late merge)
attention_registration_model.py	Transformer/attention-based registration network
transformer_model.py	Transformer encoder blocks used by the attention model
attention_layer.py	Custom attention layer building blocks
params_setting.py	Central config: choose network (RnnWalkNet / AttentionWalkNet / TransformerWalkNet), walk length, data augmentation ranges, etc.

Everything else (dataset.py, dataset_prepare.py, walks.py, utils.py, train_val.py, evaluate_segmentation.py, evaluate_classification.py, rnn_model.py) is the original MeshWalker classification/segmentation pipeline, reused here for walk generation and data loading.

Installation

A step-by-step installation guide for Ubuntu is provided in INSTALL.md.

Data

To prepare the shrec11 dataset used for registration experiments, follow the original MeshWalker data pipeline:

python dataset_prepare.py shrec11

See INSTALL.md and the original MeshWalker docs for the full dataset setup (raw dataset sources, get_datasets.sh, get_pretrained.sh).

Training

Select the network architecture (RnnWalkNet, AttentionWalkNet, or TransformerWalkNet) and walk/augmentation settings in params_setting.py, then run:

python train_object_registration.py shrec11 <part>

<part> is one of 10-10_A / 10-10_B / 10-10_C / 16-04_A / 16-04_B / 16-04_C.

Results and TensorBoard logs are written to runs/.

Evaluating

python evaluate_object_registration.py shrec11 <part> <trained model directory>

Original MeshWalker task (classification / segmentation)

The underlying MeshWalker pipeline (mesh classification and segmentation) is unmodified and can still be run via train_val.py, evaluate_segmentation.py, and evaluate_classification.py — see the original project for full instructions.

Reference

If you find this useful, please cite the original MeshWalker paper:

@article{lahav2020meshwalker,
  title={MeshWalker: Deep Mesh Understanding by Random Walks},
  author={Lahav, Alon and Tal, Ayellet},
  journal={arXiv preprint arXiv:2006.05353},
  year={2020}
}

Report

The full report, 3D Object Registration using Random Walk Algorithm, covers the architecture design space, experimental setup, and results in detail.

Questions / Issues

If you have questions or issues running this code, please open an issue.