Improved Two-Stage Neural Network for Defect Detection

Industrial defect detection can be considered as a downstream task of object detection in computer vision, and classical object detection models can be directly applied to it. However, vanilla object detect CNNs (Convolutional Neural Network), which requires a large number of samples for training, do not achieve satisfactory detection results in industrial defect detection with mostly few-shot datasets. To tackle this problem, we choose to use a shallow CNN and make some improvements to it. Specifically, first, a two-stage neural network is adopted, which consists of two sub-networks that produce defect segmentation and classification results, respectively. The network consists of convolutional and pooling layers and is trained in an end-to-end approach to perform the defect detection task. Then, on the one hand, in order to enhance the feature representation ability of network, all pooling layers are replaced with convolutional layers to implement the downsampling operation. On the other hand, in order to reduce the computation-cost and enhance the adaptability to few-shot datasets, residual connections are added to the network structure. Thus, the network to further improve the detection performance by adding a small number of parameters but without affecting the training convergence. We perform experiments on a fabric dataset containing a small number of samples and obtained the state-of-the-art (SOTA), and then perform additional ablation experiments to confirm the effectiveness of the adopted method and the improvements on a few-shot dataset.