Selective Prototype Network for Few-Shot Metal Surface Defect Segmentation

Metal Surface defects segmentation is a critical task to make pixel-level predictions about defects in the industrial production process, which has great significance in improving product quality. Existing segmentation algorithms use numerous labeled defective images for training and can not be generalized to different metal surfaces. Additionally, the metal surface has different materials and the defect samples are insufficient. That means collecting defective images and annotates pixel labels takes more time. In order to solve the above problems, this paper proposed a novel selective prototype network (SPNet) with matrix decomposition attention mechanism for few-shot metal surface defect segmentation, which aims to learn a model that can be generalized to novel surface classes with only a few labeled defect samples. Using a selective prototype acquired from the support image to learn query image, SPNet efficiently utilizes the information of the same metal surface defects and meanwhile offers sufficient representation for different metal surface defects. With this, SPNet fully utilizes correlation knowledge from the known defects and provides better generalization on unknown defects. Moreover, SPNet introduces a feature attention mechanism based on matrix decomposition. The novel attention method factorizes the complicated feature representation to acquire more accurate global context information. In addition, to improve the segmentation performance, a conditional boundary refinement module is proposed. Experimental results on the Defects dataset show that SPNet achieves state-of-the-art performance.