update version on better simplified Manifold40

.gitignore

@@ -1,2 +1,5 @@
 .idea/
-*.pkl
+ckpt_root/
+dataset/
+*.pkl
+*/__pycache__/

README.md

@@ -1,54 +1,49 @@
 ## MeshNet: Mesh Neural Network for 3D Shape Representation
 Created by Yutong Feng, Yifan Feng, Haoxuan You, Xibin Zhao, Yue Gao from Tsinghua University.
 
-![pipeline](doc/pipeline.png)
+![pipeline](doc/pipeline.PNG)
 ### Introduction
 
-This work will appear in AAAI 2019. We proposed a novel framework (MeshNet) for 3D shape representation, which could learn on mesh data directly and achieve satisfying performance compared with traditional methods based on mesh and representative methods based on other types of data. You can also check out [paper](http://gaoyue.org/paper/MeshNet.pdf) for a deeper introduction.
+This work was published in AAAI 2019. We proposed a novel framework (MeshNet) for 3D shape representation, which could learn on mesh data directly and achieve satisfying performance compared with traditional methods based on mesh and representative methods based on other types of data. You can also check out [paper](https://ojs.aaai.org/index.php/AAAI/article/view/4840/4713) for a deeper introduction.
 
 Mesh is an important and powerful type of data for 3D shapes. Due to the complexity and irregularity of mesh data, there is little effort on using mesh data for 3D shape representation in recent years. We propose a mesh neural network, named MeshNet, to learn 3D shape representation directly from mesh data. Face-unit and feature splitting are introduced to solve the complexity and irregularity problem. We have applied MeshNet in the applications of 3D shape classification and retrieval. Experimental results and comparisons with the state-of-the-art methods demonstrate that MeshNet can achieve satisfying 3D shape classification and retrieval performance, which indicates the effectiveness of the proposed method on 3D shape representation.
 
 In this repository, we release the code and data for train a Mesh Neural Network for classification and retrieval tasks on ModelNet40 dataset.
 
-### Citation
+### Update
+**[2021/12]** We have released an updated version that the proposed MeshNet achieves 92.75% classification accuracy on ModelNet40. The results are based on a better simplified version of ModelNet40, named "Manifold40", with watertight mesh and 500 faces per model. We also provide a more stable training script to achieve the performance. See the Usage section for details.
 
-if you find our work useful in your research, please consider citing:
+### Usage
 
+#### Installation
+You could install the required package as follows. This code has been tested with Python 3.8 and CUDA 11.1.
 ```
-@article{feng2018meshnet,
-  title={MeshNet: Mesh Neural Network for 3D Shape Representation},
-  author={Feng, Yutong and Feng, Yifan and You, Haoxuan and Zhao, Xibin and Gao, Yue},
-  journal={AAAI 2019},
-  year={2018}
-}
+pip install -r requirements.txt
 ```
 
-### Installation
-
-Install [PyTorch 0.4.0](https://pytorch.org). You also need to install yaml. The code has been tested with Python 3.6, PyTorch 0.4.0 and CUDA 9.0 on Ubuntu 16.04.
-
-### Usage
-
 ##### Data Preparation
-
-Firstly, you should download the [reorganized ModelNet40 dataset](https://drive.google.com/open?id=1o9pyskkKMxuomI5BWuLjCG2nSv5iePZz). Then, configure the "data_root" in `config/train_config.yaml` and `config/test_config.yaml` with your path to the downloaded dataset:
-
+MeshNet requires the pre-processed ModelNet40 with simplified and re-organized mesh data. To quickly start training, we recommend to use our [pre-processed ModelNet40 dataset](https://cloud.tsinghua.edu.cn/f/77436a9afd294a52b492/?dl=1), and configure the "data_root" in `config/train_config.yaml` and `config/test_config.yaml` with your path to the downloaded dataset. By default, run
 ```
-# config/train_config.yaml and config/test_config.yaml
-dataset:
-  data_root: [your_path_to_dataset]
+wget --content-disposition https://cloud.tsinghua.edu.cn/f/77436a9afd294a52b492/?dl=1
+mkdir dataset
+unzip -d dataset/ ModelNet40_processed.zip 
+rm ModelNet40_processed.zip 
 ```
 
-For each data file `XXX.off` in ModelNet, we reorganize it to the format required by MeshNet and store it into `XXX.npz`. The reorganized file includes two parts of data:
+The details of our pre-processing are as follows: The original dataset are from [ModelNet](http://modelnet.cs.princeton.edu/). Firstly, we simplify the mesh models with no more than `max_faces` faces. We now recommend to use the [Manifold40](https://cloud.tsinghua.edu.cn/f/2a292c598af94265a0b8/?dl=1) version with watertight mesh and `max_faces=500`. Then we reorganize the dataset to the format required by MeshNet and store it into `XXX.npz`. The reorganized file includes two parts of data:
+- The "faces" part contains the center position, vertices' positions and normal vector of each face.
+- The "neighbors" part contains the indices of neighbors of each face.
 
-* The "face" part contains the center position, vertices' positions and normal vector of each face.
-* The "neighbor_index" part contains the indices of neighbors of each face.
+If you wish to create and use your own dataset, simplify your models into `.obj` format and use the code in `data/preprocess.py` to transform them into the required `.npz` format. Notice that the parameter `max_faces` in config files should be maximum number of faces among all of your simplified mesh models. 
 
-If you wish to create and use your own dataset, simplify your models and organize the `.off` files similar to the ModelNet dataset. 
-Then use the code in `data/preprocess.py` to transform them into the required `.npz` format. 
-Notice that the parameter `max_faces` in config files should be maximum number of faces among all of your simplified mesh models. 
+##### Evaluation
+The pretrained MeshNet model weights are stored in [pretrained model](https://cloud.tsinghua.edu.cn/f/33bfdc6f103340daa86a/?dl=1). You can download it and configure the "load_model" in `config/test_config.yaml` with your path to the weight file. Then run the test script.
+```
+wget --content-disposition https://cloud.tsinghua.edu.cn/f/33bfdc6f103340daa86a/?dl=1
+python test.py
+```
 
-##### Train Model
+##### Training
 
 To train and evaluate MeshNet for classification and retrieval:
 
@@ -58,21 +53,23 @@ python train.py
 
 You can modify the configuration in the `config/train_config.yaml` for your own training, including the CUDA devices to use, the flag of data augmentation and the hyper-parameters of MeshNet.
 
-##### Test Model
-
-The pretrained MeshNet model weights are stored in [pretrained model](https://drive.google.com/open?id=1l8Ij9BODxcD1goePBskPkBcgKW76Ewcs). You can download it and configure the "load_model" in `config/test_config.yaml` with your path to the weight file.
 
-```
-# config/test_config.yaml
-load_model: [your_path_to_weight_file]
-```
+### Citation
 
-To evaluate the model for classification and retrieval:
+if you find our work useful in your research, please consider citing:
 
-```bash
-python test.py
+```
+@inproceedings{feng2019meshnet,
+  title={Meshnet: Mesh neural network for 3d shape representation},
+  author={Feng, Yutong and Feng, Yifan and You, Haoxuan and Zhao, Xibin and Gao, Yue},
+  booktitle={Proceedings of the AAAI Conference on Artificial Intelligence},
+  volume={33},
+  number={01},
+  pages={8279--8286},
+  year={2019}
+}
 ```
 
 ### Licence
 
-Our code is released under MIT License (see LICENSE file for details).
+Our code is released under MIT License (see LICENSE file for details).

config/config.py

@@ -14,7 +14,7 @@ def _check_dir(dir, make_dir=True):
 
 def get_train_config(config_file='config/train_config.yaml'):
     with open(config_file, 'r') as f:
-        cfg = yaml.load(f)
+        cfg = yaml.load(f, Loader=yaml.loader.SafeLoader)
 
     _check_dir(cfg['dataset']['data_root'], make_dir=False)
     _check_dir(cfg['ckpt_root'])
@@ -24,7 +24,7 @@ def get_train_config(config_file='config/train_config.yaml'):
 
 def get_test_config(config_file='config/test_config.yaml'):
     with open(config_file, 'r') as f:
-        cfg = yaml.load(f)
+        cfg = yaml.load(f, Loader=yaml.loader.SafeLoader)
 
     _check_dir(cfg['dataset']['data_root'], make_dir=False)
 

config/test_config.yaml

@@ -3,12 +3,12 @@ cuda_devices: '0'
 
 # dataset
 dataset:
-  data_root: 'ModelNet40_MeshNet/'
+  data_root: 'dataset/ModelNet40_processed'
   augment_data: false
-  max_faces: 1024
+  max_faces: 500
 
 # model
-load_model: 'MeshNet_best_9192.pkl'
+load_model: 'MeshNet_ModelNet40_250e_bs128_lr6e-4.pkl'
 
 # MeshNet
 MeshNet:
@@ -17,3 +17,10 @@ MeshNet:
     sigma: 0.2
   mesh_convolution:
     aggregation_method: 'Concat'  # Concat/Max/Average
+  mask_ratio: 0.95
+  dropout: 0.5
+  num_classes: 40
+
+# test config
+batch_size: 128
+retrieval_on: true

config/train_config.yaml

@@ -1,14 +1,16 @@
 # CUDA
-cuda_devices: '0'   # multi-gpu training is available
+cuda_devices: '0,1'   # multi-gpu training is available
 
 # dataset
 dataset:
-  data_root: 'ModelNet40_MeshNet/'
+  data_root: 'dataset/ModelNet40_processed'
+  max_faces: 500
   augment_data: true
-  max_faces: 1024
+  jitter_sigma: 0.01
+  jitter_clip: 0.05
 
 # result
-ckpt_root: 'ckpt_root/'
+ckpt_root: 'ckpt_root'
 
 # MeshNet
 MeshNet:
@@ -17,12 +19,20 @@ MeshNet:
     sigma: 0.2
   mesh_convolution:
     aggregation_method: 'Concat'  # Concat/Max/Average
+  mask_ratio: 0.95
+  dropout: 0.5
+  num_classes: 40
 
 # train
-lr: 0.01
+seed: 0
+lr: 0.0006
 momentum: 0.9
 weight_decay: 0.0005
-batch_size: 64
-max_epoch: 150
-milestones: [30, 60]
+batch_size: 128
+max_epoch: 250
+optimizer: 'adamw' # sgd/adamw
+scheduler: 'cos' # step/cos
+milestones: [30, 60, 90]
 gamma: 0.1
+retrieval_on: true   # enable evaluating retrieval performance during training
+save_steps: 10

data/ModelNet40.py

@@ -19,12 +19,17 @@ class ModelNet40(data.Dataset):
 
     def __init__(self, cfg, part='train'):
         self.root = cfg['data_root']
-        self.augment_data = cfg['augment_data']
         self.max_faces = cfg['max_faces']
         self.part = part
+        self.augment_data = cfg['augment_data']
+        if self.augment_data:
+            self.jitter_sigma = cfg['jitter_sigma']
+            self.jitter_clip = cfg['jitter_clip']
 
         self.data = []
         for type in os.listdir(self.root):
+            if type not in type_to_index_map.keys():
+                continue
             type_index = type_to_index_map[type]
             type_root = os.path.join(os.path.join(self.root, type), part)
             for filename in os.listdir(type_root):
@@ -34,14 +39,14 @@ def __init__(self, cfg, part='train'):
     def __getitem__(self, i):
         path, type = self.data[i]
         data = np.load(path)
-        face = data['face']
-        neighbor_index = data['neighbor_index']
+        face = data['faces']
+        neighbor_index = data['neighbors']
 
         # data augmentation
         if self.augment_data and self.part == 'train':
-            sigma, clip = 0.01, 0.05
-            jittered_data = np.clip(sigma * np.random.randn(*face[:, :12].shape), -1 * clip, clip)
-            face = np.concatenate((face[:, :12] + jittered_data, face[:, 12:]), 1)
+            # jitter
+            jittered_data = np.clip(self.jitter_sigma * np.random.randn(*face[:, :3].shape), -1 * self.jitter_clip, self.jitter_clip)
+            face = np.concatenate((face[:, :3] + jittered_data, face[:, 3:]), 1)
 
         # fill for n < max_faces with randomly picked faces
         num_point = len(face)