first commit

README.md

@@ -0,0 +1,74 @@
+## 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)
+### 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.
+
+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
+
+if you find our work useful in your research, please consider citing:
+
+```
+@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}
+}
+```
+
+### 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=1l8Ij9BODxcD1goePBskPkBcgKW76Ewcs). Then, configure the "data_root" in `config/train_config.yaml` and `config/test_config.yaml` with your path to the downloaded dataset:
+
+```yaml
+# config/train_config.yaml and config/test_config.yaml
+dataset:
+	data_root: [your_path_to_dataset]
+```
+
+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 "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.
+
+##### Train Model
+
+To train and evaluate MeshNet for classification and retrieval:
+
+```bash
+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=1m5Uy9-oXMNPZ129owKvQ5ipH3f0vdABs). You can download it and configure the "load_model" in `config/test_config.yaml` with your path to the weight file.
+
+```yaml
+# config/test_config.yaml
+load_model: [your_path_to_weight_file]
+```
+
+To evaluate the model for classification and retrieval:
+
+```bash
+python test.py
+```
+
+### Licence
+
+Our code is released under MIT License (see LICENSE file for details).

config/__init__.py

@@ -0,0 +1 @@
+from .config import get_train_config, get_test_config

config/config.py

@@ -0,0 +1,31 @@
+import os
+import os.path as osp
+import yaml
+
+
+def _check_dir(dir, make_dir=True):
+    if not osp.exists(dir):
+        if make_dir:
+            print('Create directory {}'.format(dir))
+            os.mkdir(dir)
+        else:
+            raise Exception('Directory not exist: {}'.format(dir))
+
+
+def get_train_config(config_file='config/train_config.yaml'):
+    with open(config_file, 'r') as f:
+        cfg = yaml.load(f)
+
+    _check_dir(cfg['dataset']['data_root'], make_dir=False)
+    _check_dir(cfg['ckpt_root'])
+
+    return cfg
+
+
+def get_test_config(config_file='config/test_config.yaml'):
+    with open(config_file, 'r') as f:
+        cfg = yaml.load(f)
+
+    _check_dir(cfg['dataset']['data_root'], make_dir=False)
+
+    return cfg

config/test_config.yaml

@@ -0,0 +1,18 @@
+# CUDA
+cuda_devices: '0'
+
+# dataset
+dataset:
+  data_root: 'ModelNet40_MeshNet/'
+  augment_data: false
+
+# model
+load_model: 'MeshNet_best_9192.pkl'
+
+# MeshNet
+MeshNet:
+  structural_descriptor:
+    num_kernel: 64
+    sigma: 0.2
+  mesh_convolution:
+    aggregation_method: 'Concat'  # Concat/Max/Average

config/train_config.yaml

@@ -0,0 +1,27 @@
+# CUDA
+cuda_devices: '0'   # multi-gpu training is available
+
+# dataset
+dataset:
+  data_root: 'ModelNet40_MeshNet/'
+  augment_data: true
+
+# result
+ckpt_root: 'ckpt_root/'
+
+# MeshNet
+MeshNet:
+  structural_descriptor:
+    num_kernel: 64
+    sigma: 0.2
+  mesh_convolution:
+    aggregation_method: 'Concat'  # Concat/Max/Average
+
+# train
+lr: 0.01
+momentum: 0.9
+weight_decay: 0.0005
+batch_size: 64
+max_epoch: 150
+milestones: [30, 60]
+gamma: 0.1

data/ModelNet40.py

@@ -0,0 +1,60 @@
+import numpy as np
+import os
+import torch
+import torch.utils.data as data
+
+
+class ModelNet40(data.Dataset):
+
+    def __init__(self, cfg, part='train'):
+        self.root = cfg['data_root']
+        self.augment_data = cfg['augment_data']
+        self.part = part
+
+        self.data = []
+        type_index = 0
+        for type in os.listdir(self.root):
+            type_root = os.path.join(os.path.join(self.root, type), part)
+            for filename in os.listdir(type_root):
+                if filename.endswith('.npz'):
+                    self.data.append((os.path.join(type_root, filename), type_index))
+            type_index += 1
+
+    def __getitem__(self, i):
+        path, type = self.data[i]
+        data = np.load(path)
+        face = data['face']
+        neighbor_index = data['neighbor_index']
+
+        # 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)
+
+        # fill for n < 1024
+        num_point = len(face)
+        if num_point < 1024:
+            fill_face = []
+            fill_neighbor_index = []
+            for i in range(1024 - num_point):
+                index = np.random.randint(0, num_point)
+                fill_face.append(face[index])
+                fill_neighbor_index.append(neighbor_index[index])
+            face = np.concatenate((face, np.array(fill_face)))
+            neighbor_index = np.concatenate((neighbor_index, np.array(fill_neighbor_index)))
+
+        # to tensor
+        face = torch.from_numpy(face).float()
+        neighbor_index = torch.from_numpy(neighbor_index).long()
+        target = torch.tensor(type, dtype=torch.long)
+
+        # reorganize
+        face = face.permute(1, 0).contiguous()
+        centers, corners, normals = face[:3], face[3:12], face[12:]
+        corners = corners - torch.cat([centers, centers, centers], 0)
+
+        return centers, corners, normals, neighbor_index, target
+
+    def __len__(self):
+        return len(self.data)

data/__init__.py

@@ -0,0 +1 @@
+from .ModelNet40 import ModelNet40

models/MeshNet.py

@@ -0,0 +1,53 @@
+import torch
+import torch.nn as nn
+from models import SpatialDescriptor, StructuralDescriptor, MeshConvolution
+
+
+class MeshNet(nn.Module):
+
+    def __init__(self, cfg, require_fea=False):
+        super(MeshNet, self).__init__()
+        self.require_fea = require_fea
+
+        self.spatial_descriptor = SpatialDescriptor()
+        self.structural_descriptor = StructuralDescriptor(cfg['structural_descriptor'])
+        self.mesh_conv1 = MeshConvolution(cfg['mesh_convolution'], 64, 131, 256, 256)
+        self.mesh_conv2 = MeshConvolution(cfg['mesh_convolution'], 256, 256, 512, 512)
+        self.fusion_mlp = nn.Sequential(
+            nn.Conv1d(1024, 1024, 1),
+            nn.BatchNorm1d(1024),
+            nn.ReLU(),
+        )
+        self.concat_mlp = nn.Sequential(
+            nn.Conv1d(1792, 1024, 1),
+            nn.BatchNorm1d(1024),
+            nn.ReLU(),
+        )
+        self.classifier = nn.Sequential(
+            nn.Linear(1024, 512),
+            nn.ReLU(),
+            nn.Dropout(p=0.5),
+            nn.Linear(512, 256),
+            nn.ReLU(),
+            nn.Dropout(p=0.5),
+            nn.Linear(256, 40)
+        )
+
+    def forward(self, centers, corners, normals, neighbor_index):
+        spatial_fea0 = self.spatial_descriptor(centers)
+        structural_fea0 = self.structural_descriptor(corners, normals, neighbor_index)
+
+        spatial_fea1, structural_fea1 = self.mesh_conv1(spatial_fea0, structural_fea0, neighbor_index)
+        spatial_fea2, structural_fea2 = self.mesh_conv2(spatial_fea1, structural_fea1, neighbor_index)
+        spatial_fea3 = self.fusion_mlp(torch.cat([spatial_fea2, structural_fea2], 1))
+
+        fea = self.concat_mlp(torch.cat([spatial_fea1, spatial_fea2, spatial_fea3], 1))
+        fea = torch.max(fea, dim=2)[0]
+        fea = fea.reshape(fea.size(0), -1)
+        fea = self.classifier[:-1](fea)
+        cls = self.classifier[-1:](fea)
+
+        if self.require_fea:
+            return cls, fea / torch.norm(fea)
+        else:
+            return cls

models/__init__.py

@@ -0,0 +1,2 @@
+from .layers import SpatialDescriptor, StructuralDescriptor, MeshConvolution
+from .MeshNet import MeshNet