# Defect Detection on Rolling Element Surface Scans Using Neural Image Segmentation
## 5 Classification  
### Neural Network Architecture  
As stated in Section 3, we propose the use of a neural segmentation network to localize defects on the input images. Image segmentation problems are solved in the literature with convolutional encoder-decoder architectures [11]. These remap the features extracted by a traditional CNN trunk, back into an image of the input size. The use of an existing pre-trained network was foregone. Architectures like SegNet [11] are usually trained on segmentation tasks for scene understanding problems. The structures found on the steel surface scans of rolling elements bear little resemblance to the structures found in these scenes. This results in a suboptimal feature extraction, when such pre-trained networks are used. Additionally, existing architectures are usually quite large with 26 and 23 layers of convolutions for SegNet and U-Net [12], respectively. This depth is required for their complex task of segmenting images into dozens of classes or finding segmentations in complex biomedical images. Instead, a custom designed architecture as listed in Table 1 was used. It is of smaller depth, as the given defect detection problem is a more specific and therefore less complex problem. Furthermore, a smaller architecture allows for shorter training time and results in faster inference time during deployment. The encoder CNN consists of three stacked layers of convolutions. Each layer uses filter kernels of size \(7\times 7\) pixels. After all convolution layers, Leaky-ReLU activations are used as non-linear activation functions. The first layer convolves the image with 16 filters. Every subsequent layer in the encoder doubles the filter count. A stride of two in all three convolution layers reduces the size of the activation images in both dimensions by a total factor of eight. A reduction of the activation map sizes through the more common max pooling operation was forgone, based on the conclusions drawn in [13]. The bottleneck resulting from the size reductions reduces the representational capability of the network and acts regularizing. Since every pixel in a segmentation problem is equally important, the convolution operations also act on edge pixels of the image. To allow this, reflection padding extends the size of the image before convolutions are applied. Most CNN architectures use zero-padding to fill the space around the image. In our experiments, this led to artefacts on the edges of activation maps. The padded zeros can be interpreted as defect image regions and result in erroneous behavior. To combat this problem, reflection padding extends the image by reflecting inner image pixels outwards. \begin{table}
\begin{tabular}{c c c c c} \hline \hline  & **Layers** & **Kernel Size** & **Filters** & **Stride** \\ \cline{2-5}
**Encoder** & Conv1 & \(7\times 7\) & 16 & /2 \\  & Conv2 & \(7\times 7\) & 32 & /2 \\  & Conv3 & \(7\times 7\) & 64 & /2 \\  & Resize1 & - & - & \(\times 2\) \\ \hline  & Conv4 & \(7\times 7\) & 32 & 1 \\  & Resize2 & - & - & \(\times 2\) \\  & Conv5 & \(7\times 7\) & 16 & 1 \\  & Resize3 & - & - & \(\times 2\) \\  & Conv6 & \(7\times 7\) & 8 & 1 \\  & Conv7 & \(1\times 1\) & 1 & 1 \\ \hline \hline \end{tabular}
\end{table} 
Table 1: Architecture of the utilized neural segmentation network. These padded pixel values are more plausible and reduce unwanted detection in the activation map, as is apparent in Figure 4.  
To regularize the network, spatial dropout instead of regular dropout is added after every convolution layer in the encoder. Regular dropout is less effective in convolution networks, as is discussed in [14]. The authors of [14] use a 10% drop rate in their CNN. Our experiments showed that further regularization at 30% drop rate results in best performance, given the low amount of training data. The decoder part of the segmentation network decodes the latent representation into a segmentation mask. It stacks three resize convolutions, starting with 32 kernels. Every subsequent layer halves the filter count.