This is the PyTorch implementation of our MeshMAE.
- python 3.9+
- CUDA 11.1+
- torch 1.11+
To install python requirements:
pip install -r requirements.txt
To install PyGem, please refer to Pygem.
Here, we provide the download links of the datasets for pre-train, classification and segmentation.
- ModelNet40 here
- Humanbody here
- COSEG-aliens here
- ShapeNet here(we also provide the processed ShapeNet dataset as here)
Please create a new folder 'datasets' in the main root, and put the downloaded datasets in this folder. And '--dataroot' in the 'xxx.sh' refers to the root of datasets.
For example, the root of ModelNet40 should be:
--dataroot ./dataset/Manifold40-MAPS-96-3/
To process the raw data, please use data_preprocess/manifold.py, which can transform non-manifold mesh data into manifold and simplify it to 500 faces.
To remesh the meshes, you can refer to the datagen_maps.py of SubdivNet, which can generate hierarchical structures.
Here, we provide the download links of the pre-trained models.
Task | Dataset | Baseline | Finetune |
---|---|---|---|
Pre-training | ShapeNet | here | --- |
Classification | ModelNet40 | here | here |
Segmentation | HumanBody | here | here |
Segmentation | Coseg-alien | here | here |
- To pretrain on the ShapeNet dataset, run this command:
sh scripts/pretrain/train_pretrain.sh
- To train the classification model from scratch, run this command:
sh scripts/classification/train.sh
- To finetune the classification model, run this command:
sh scripts/classification/train_finetune.sh
- To train the segmentation model from scratch, run this command:
sh scripts/segmentation/<...>/train.sh
- To finetune the segmentation model, run this command:
sh scripts/segmentation/<...>/train_finetune.sh
To finetune the model, please create a folder 'checkpoints' in the main root, and put the the pre-trained model in it. And '--checkpoint' in the 'train_finetune.sh' refers to the root the pre-trained model.
For example, the root of pre-trained model should be:
--checkpoint "./checkpoints/shapenet_pretrain.pkl"
@inproceedings{meshmae2022,
title={MeshMAE: Masked Autoencoders for 3D Mesh Data Analysis},
author={Liang, Yaqian and Zhao, Shanshan and Yu, Baosheng and Zhang, Jing and He, Fazhi},
booktitle={European Conference on Computer Vision},
year={2022},
}