Attention-Guided Multitask Learning for Surface Defect Identification
Vignesh Sampath, Iñaki Maurtua, Juan José Aguilar Martín, Andoni Rivera, Jorge Molina, Aitor Gutierrez
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.