Attention-Guided Multitask Learning for Surface Defect Identification

Abstract

Surface defect identification is an essential task in the industrial quality control process, in which visual checks are conducted on a manufactured product to ensure that it meets quality standards. The convolutional neural network (CNN)-based surface defect identification method has proven to outperform traditional image processing techniques. However, the real-world surface defect datasets are limited in size due to the expensive data generation process and the rare occurrence of defects. To address this issue, this article presents a method for exploiting auxiliary information beyond the primary labels to improve the generalization ability of surface defect identification tasks. Considering the correlation between pixel-level segmentation masks, object-level bounding boxes, and global image-level classification labels, we argue that jointly learning features of the related tasks can improve the performance of surface defect identification tasks. This article proposes a framework named Defect-Aux-Net, based on multitask learning with attention mechanisms that exploit the rich additional information from related tasks with the goal of simultaneously improving robustness and accuracy of the CNN-based surface defect identification. We conducted a series of experiments with the proposed framework. The experimental results showed that the proposed method can significantly improve the performance of state-of-the-art models while achieving an overall accuracy of 97.1%, Dice score of 0.926, and mean average precision of 0.762 on defect classification, segmentation, and detection tasks.