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. Pygem不能直接pip装 pip装的版本太老 这里参考:https://blog.csdn.net/python1639er/article/details/123172311
关于chamfer这个包,下载代码有自带的.cpp,需要手动安装:
先控制台 cd chamfer_dist
然后dir
再 Python setup.py build
没问题的话 再执行
python setup.py install
再次运行代码发现这个包不再报错
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.
更改后的dataset.py读取数据的存放路径格式为:
| datasets_name
├── train/
│ ├── airplane_0001.obj
│ ├── airplane_0002.obj
├── test/
│ ├── airplane_0001.obj
│ ├── airplane_0002.obj
便于直接处理训练与测试数据分开存放的数据
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},
}
如需对预训练进行test,在meshmae.py倒数几行,需要更改return值