3D Object Classification and Segmentation with MeshCNN and PyTorch
By: Mason Mcgough
Tags: Machine Learning, 3d, Deep Learning, Pytorch, Meshcnn

3D Object Classification and Segmentation with MeshCNN and PyTorch

Successive mesh pooling operations applied to a 3D model of a centaur [From the MeshCNN project page].

Applying machine learning to 3D data is not as straightforward as it is for images. There are many representations for 3D data, but the choice dictates what kinds of learning strategies you can employ. In this article I cover a particularly interesting (at least to me 😆) strategy called MeshCNN: A Network with an Edge. This paper describes a generic framework for processing 3D models for classification and segmentation tasks. Perhaps its most interesting feature is its mesh pooling operation, which enables us to combine features of a mesh at multiple scales (similar to a vision CNN). It is a learned operation that gradually reduces a model to the edges that are most informative to a given task. MeshCNN combines many of the best properties of each popular 3D representation. Before we go into details however, let us go over what those properties are with a brief review of 3D representations.

3D Data Representations

What is the best way to represent 3D meshes for deep learning? Unlike 2D RGB images, there is no consensus on the best representation. The question is a difficult one to answer because the choice of representation determines the learning approach we must take. For the classification example, you can project the model from 3D space into 2D images and apply standard 2D convolutions. You can represent the 3D space occupied by the model as a grid of voxels, allowing you to apply 3D convolutions. You can simply sample the vertices of the mesh as a 3D point cloud and apply specialized methods like PointNet++ or 3D point capsule networks. There are even methods like PolyGen that process the vertices and faces of the model directly, which I touch on in my previous article Generating 3D Models with PolyGen and PyTorch.

The most typical representations for 3D meshes [From the Kaolin library by NVIDIA, licensed under the Apache License Ver. 2.0 and edited to add text].

Projection

Some of the earliest 3D deep learning studies bypass the 3D representation issue directly and simply project the 3D model into 2D images. This makes 3D data amenable to the classic vision CNN methods. The method described in Multi-view Convolutional Neural Networks for 3D Shape Recognition for instance projects a model into 12 unique viewpoints and pools their activations to produce a softmax score. On the ModelNet40 dataset they report a solid classification accuracy of 90.1%. The best performance as of today on the ModelNet40 dataset according to the project’s website is 97.37%, which comes from a paper called RotationNet: Joint Object Categorization and Pose Estimation Using Multiviews from Unsupervised Viewpoints. Similar to the previous paper, it trains on multiple views to predict an object category, but simultaneously predicts the viewpoint as well, treating it as a latent variable. It is able to effectively predict both object class and viewpoint, even for real-world objects.

While simple and elegant, the projection representation a) does not take the model’s full topology into account, b) makes assumptions about how the model should be viewed, and c) does not provide a straightforward solution for non-global tasks like segmentation.

Voxel

The voxel method addresses many of the issues with the projection method without sacrificing convolutions. Create a dense occupancy grid by dividing the 3D space into bins, assigning a boolean value to each cell if it falls inside of the model. The grid layout is simple to apply a 3D convolution to. Possibly the earliest paper to use this representation is 3D ShapeNets: A Deep Representation for Volumetric Shapes, the paper which actually introduces the ModelNet40 dataset. In this study, they attempt to recover the 3D structure of an object by predicting 3D voxels from a single depth image and achieve pretty remarkable results.

Classification from 3D voxels is another story. While you can achieve decent results, the downside is excessive amounts of convolutions on empty space, since the occupancy of the space is highly sparse. On top of that, model resolution scales the number of weights cubically, so sizes greater than 256x256x256 in most cases are impractical, even when significantly exploiting sparsity of the space, such as in OctNet: Learning Deep 3D Representations at High Resolutions. In this paper they achieve about 86% accuracy on ModelNet40 with great improvements to speed and memory consumption. High-resolution voxel spaces are costly, but low-resolution voxel spaces squash the potentially useful fine details of model topology.

Point Cloud

Many methods focus on processing vector point clouds directly. For instance, in the 2019 paper Spherical Kernel for Efficient Graph Convolution on 3D Point Clouds the authors achieve 89.3% accuracy on the ModelNet40 task using their method. The benefit of the 3D point cloud representation is its versatility, since everything from LiDAR scans to authored 3D models can be represented as a 3D point cloud. Even the classic PointNet and PointNet++ models can achieve pretty good results on the classification task (88.0% in the paper above). The downside is that most point cloud methods disregard surfaces entirely. Without faces, the mesh’s true topology cannot be known, since any set of vectors does not uniquely define a set of faces.

The knees of this camel are so close that they may not be separable without knowledge of the model surface, a.k.a. its faces[From Fig. 3 of the MeshCNN paper].

One of the best-performing methods on ModelNet40 comes from a method called Relation-Shape Convolutional Neural Network for Point Cloud Analysis which achieves 93.6% accuracy. This method, named RS-CNN, attempts to infer the underlying topology of a given point cloud from geometric priors, granting the model spatial awareness of its input points. This model achieves excellent performance and can be applied to point clouds as well as meshes. However, it has no mechanism to utilize mesh information even when it is available.

MeshCNN

Is there a way to study the mesh directly, without sacrificing valuable topology information, suffering the computational cost of voxels, or making assumptions about how to view it? The solution proposed by MeshCNN is to study the vertices and their connections, or edges, jointly by treating the 3D model as a graph or manifold. This method defines convolution and pooling layers over the edges of 3D meshes, allowing us to more or less use the standard toolset of convolutional neural nets. Ultimately they are able to achieve 98.6% accuracy on 30 classes from the SHREC 11 dataset (they do not report accuracy for ModelNet40) as well as impressive segmentation performance on object parts and human body datasets. The MeshCNN source code is provided by the authors if you wish to refer to their implementation.

A depiction of the mesh pooling operation [From Fig. 2 of the MeshCNN paper]. a) A given edge (red) of a triangle mesh has exactly 4 neighbors (blue). b) The pooling operation dissolves this edge by merging its two vertices, which in turn merges the edge pairs on either side of the dissolved edge, c) resulting in 2 edges.

The edge of any watertight 3D mesh is incident on exactly two faces (with the exception of boundary or non-manifold edges). If the mesh is triangular (i.e. all of its faces have exactly 3 edges), then the two incident faces connected to any edge will always touch a total of 5 edges. This consistency makes triangular meshes particularly convenient for machine learning techniques. MeshCNN makes the assumption that our models are all triangular, manifold meshes for this reason. Typical 3D reconstruction methods like photogrammetry, structured light scanning, and laser scanning produce triangle meshes, so this method may be directly applicable to such data. Authored meshes, which typically contain quads or n-gons, can usually be converted to triangle meshes with a few clicks in Maya or Blender.

If we can define a convolution over an edge and its 4 neighbors, we can subsequently build a convolutional neural network that processes the whole mesh. The challenge becomes defining a set of operations that have the following properties:

Must be invariant to vertex or edge order (locally or globally).

Must be invariant to similarity transformations (i.e. mesh translation, rotation, and scale).

Must convey the relationship of a given vertex or edge to its neighbors and its place in the global structure of the mesh.

Mesh Convolution

Edge and its neighbors (left) [From Fig. 4 of the MeshCNN paper]. The edges a,c (red) and b,d (blue) are opposing pairs. Note the clockwise symbols indicating the strict ordering of edge pairs such that a always precedes b and c precedes d. The input features for a given edge (right) are also designed to be invariant to ordering of edges.

First let’s look at the convolution operation itself. Given an edge and 4 neighbors, each of which has its own features, the convolution needs to be invariant to the ordering of these edges. The simple solution employed in the paper is to define the convolution using symmetric functions. They process pairs of opposing edges (e.g. a,c and b,d), ordered clockwise with respect to the center edge e and take their finite sum and differences before applying convolution:

Formula for the equivariant feature combination [Eq. 2 from the MeshCNN paper].

Code for equivariant convolution.

Mesh Input Features

However we haven’t yet solved a critical question. What features do we start with? We must create an input feature representation that is analogous to the RGB channels in 2D images before we can apply our first convolution. For that the authors simply define the dihedral angle (angle between two adjoining faces), the symmetric opposite angles (the angle of the opposing corners, sorted to maintain order invariance), and two edge-length ratios (height/base ratio for each triangle, also sorted) for each face for a total of 5 input features.

Code for input features.

Mesh Pooling

Lastly let’s look at the pooling operation. Pooling simply merges the two vertices of the edge together, collapsing the 5 edges of a given neighborhood into 2. The features of the two new edges are simply the average of the features of the edges that they used to be. For instance, the features of our two new edges in the previous figure would be avg(a,b,e) and avg(c,d,e).

But how do we decide which edges to collapse? Every mesh pooling layer is instantiated with a target number of edges to keep (using the --pool_res argument in the code). The mesh pooling layer simply sorts the edges by the squared magnitude of the edge features. It then iteratively collapses the mesh edges until the target number of edges is achieved.

Mesh pooling learns to preserve the edges that are most informative to a given task [From Fig. 1 of the MeshCNN paper]. Pooling for the “handle or no handle” task (top) preserves the handles in the vase whereas for the “neck or no neck” task it retains the neck of the vase.

Since the mesh pooling operation is learned, it grants the model the freedom to learn weights that optimize for a given task. There is also a mesh unpooling operation that reverts the pooling, which is essential for the segmentation task. That means the network must keep track of the pooling operations done in the encoder stage of the U-Net style segmentation network. MeshCNN does this by tracking the edge collapse operations in the history_data attribute of the mesh.

Architecture

The network architecture is composed of MResConv layers, each of which consists of an initial mesh convolution (MeshConv), several successive ReLU+BatchNorm+MeshConv layers, and a residual connection followed by another ReLU. The network follows a pattern of MResConv +Norm+ MeshPool several times before ending in the task layer. For classification the task layer is simple global average pooling followed by two fully connected layers. The segmentation network is a U-Net style encoder-decoder.

Conclusion

The equivariant convolution operation, the invariant input features, and the learned mesh pooling operation make MeshCNN a particularly interesting model with the following key benefits:

More efficient and fewer parameters than older methods.

Utilizes topology of mesh (i.e. vertex and face information) rather than treating it as a point cloud.

Mesh convolution retains the convenient properties of convolutions but allows application over graph data. The 5 input features of the 3D mesh are analogous to the RGB features of an input image.

Invariant to rotation, translation, and scaling (see Section 5.4 of paper).

Mesh pooling (i.e. learned edge collapse) allows the network to learn task-specific pooling by collapsing 5 edges into 2 and dissolving both faces.

If you are eager to try it yourself, I encourage you to check out the MeshCNN source code provided by the authors. It is written in PyTorch and provides a number of helpful scripts to try the model out on the various datasets discussed in the paper. I found it very clean and easy to use, so it is worth taking a look if you are curious to look under the hood. As always, thanks for reading and keep on learning!