Finding and classifying parts by shape

Exploring the potential of machine learning in engineering design.

Introduction

In engineering design leading to the development of new parts, we can potentially speed up the process and reduce costs by starting from similar parts.

How can we find similar part models automatically? What does it mean for two objects to be similar in the first place?

We tackle the first question using machine learning on a large dataset of CAD models to find a mapping from CAD models to a low-dimensional vector space (n-tuples with the usual dot product and Euclidean distance - henceforth referred to as a latent space) such that similar models are close by. Finding similar parts is then achieved by simply returning parts that are close by in the latent space.

We also show how the latent space can be used to classify parts given only a small set of examples. The hard work of discovering the shape features that enable classification is implicit in the mapping into the latent space.

These capabilities underpin the functionality of our design tools.

Training Data

To learn a mapping from CAD models into a latent space, we require training data. We focus on datasets that use the neutral STEP file format as this provides a structural representation that makes explicit the geometry of faces and their topological connectivity. This contrasts with triangle-based formats such as STL, where the relational geometry at a structural level is less readily available. The STEP file format is also widely used in publicly available datasets. One of the largest is the ABC-Dataset of around one millions (uncategorised) models. In our work, we use the FabWave dataset. This contains around 4,500 CAD models, each labelled with one of 43 part categories. These category labels allow a simple measure of similarity to be defined, as will be seen below.

Representing CAD Models as Graphs

The topological and geometric shape information contained in a CAD model is naturally represented as a graph structure. As the first stage of mapping to a latent space, we convert each CAD model into a graph, where each node represents a face of the model and each edge represents the connection between a pair of faces that are in contact with one another. Each node has a single attribute, which is the type of the corresponding face (e.g., plane, cylinder) as defined in the Step-File ISO 10303 standard. Each edge has three attributes: (1) the angle between the connected surfaces, (2) the ratio of the sizes of the two surfaces, and (3) the type of the curve defined by the intersection of the two connected surfaces (e.g., line, circle, ellipse). Each edge is actually a pair of directed edges, one in each direction. Thus, our graph is formally an attributed directed graph. We refer to the graphs as model graphs.

CAD models represented as model graphs

Graph Neural Network

We map model graphs into a low-dimensional vector space using a neural network made up of a Graph Neural Network (GNN), global mean pooling of the node attributes and a multi-layer fully connected network. The GNN stage updates node attributes as a parameterised function of the attributes of the neighbouring nodes and connecting edges. This update is repeated multiple times so that the attributes at individual nodes become influenced by node and edge attributes from further and further away in the graph structure. Each iteration uses a different set of function parameters and is referred to as a layer of the GNN. Finally, we take the mean of all node attributes and pass this through a multi-layer fully connected neural network (two layers in our experiments), giving the final latent vector.

The initial node attributes are binary vectors of size 9, representing the nine possible face types in so-called 1-hot form. After a single layer, the node attributes become size 32 and remain at this size through further layers. The edge attributes are size 11 and remain fixed throughout.

In the diagram below, the arrays passing between layers encode the attribute vectors of a graph with N nodes and M edges.

Network Training

To train our network, we use the FabWave dataset of CAD models. We train the network end-to-end in the sense that all network parameters are estimated together (the parameter values in each GNN layer, and the parameters of the final fully-connected network). We use a triplet loss function, which compares computed distances with ground-truth for an ‘anchor’ part with a similar part and with a dissimilar part. This penalises incorrect ordering of pairwise similarities. Training the network took around 2 hours using a single Nvidia V100 GPU. In training, two parts are similar if they are from the same category and disimilar otherwise.

Results

To evaluate performance of our network we held-out a test set (20%) from FabWave and measured prediction accuracy. Specifically, we predict which of two model parts B, C is in the same class as a third model part A, based on the relative Euclidean distance between A&B and A&C in the latent space. We vary the number of GNN layers between 0 and 4 to motivate our choice of three layers in further experiments. The maximum accuracy of 98.7% is obtained with three layers. This gives a mapping of model graphs into a latent space optimised to minimise a triplet loss on the FabWave dataset.

Applications

Part Find

Here we demonstrate how similar parts cluster together in the latent space. A subset of 71 parts selected from a case study (protype sintering machine) are projected onto the first two principal components of the distribution of the 71 parts in the 16-D latent space. The part models are rendered in their approximate positions within this 2D space. In this scatterplot, an emergent clustering of parts can be seen to the human eye. Three such clusters labelled A, B and C have been highlighted. Cluster A contains circular ring-like objects. Cluster B contains hinge-like assemblies. Cluster C has a number of flat triangular pieces. Other clusters are also apparent.

The practical application of this is that given one part mapped into the latent space, a search can be carried out for the closest vectors in that space to find similar parts.

This is an interactive plot of the latent space

Manufacturing Method Prediction

The challenge posed in the case study is to determine the manufacturing process for any model part. We are given a small dataset of labels for the putative manufacturing process for just 71 parts. Our approach is to use the latent space as an intermediate representation for parts and to train a conventional fully connected neural network with three hidden layers on this space, instead of attempting to train a classifier on the original model graphs. Here we are leveraging the mapping into a low-dimensional latent space that has been optimised on a large dataset of part models (from FabWave). This is an instance of transfer learning.

Network architecture of manufacturing process classifier

The classifier is trained on 90% of the data (71 models) over 100 epochs and tested on the remaining 10% (7 models). We use 10-fold cross-validation and obtain an estimate of 70.53% for classification accuracy with a standard deviation of 6.69.

To visualise the input to the classifier, Figure 7 shows a projection of the labelled 71 parts from the sintering machine onto the first two principal components of their distribution in the latent space. The colour of points denotes the manufacturing process. The partial clustering of part models into manufacturing process is already apparent to the human eye. This intermediate representation of a latent space means that much less training data is required for the classifier.

All 71 parts of the sintering machine projected onto the first two principal components of the part distribution in the latent space. The parts are colour-coded to represent the manufacturer process

Further Resources

For a more in depth look at the machine learning aspects of the project, a paper will soon be published and linked here. The related code, which runs using Python can be found on GitHub: github.com/thazlehurst/partfind

The code respository includes a pre-trained model.