RDPrompter: Reference-Defect Prompt Learning for Few-Shot Defect Segmentation Based on Visual Foundation Model

Few-shot industrial defect segmentation (FIDS) is an extremely challenging task in industrial inspection, which focuses on segmenting unseen defect categories with only a few samples. Existing few-shot segmentation methods are commonly developed with constrained feature extraction capabilities, making it difficult to accurately segment diverse unseen defects from distinctive industrial scenarios. To this end, we aim to leverage strong generalization strengths of the large-scale visual foundation model, segment anything model (SAM), to handle FIDS task. However, simply incorporating SAM into FIDS is insufficient for automatically categorizing defects, as the SAM heavily relies on manual prompts for segmentation. To address this, we propose a novel reference-defect prompt learning method RDPrompter, which learns adaptive prompts from defect similarities based on SAM for guiding automated FIDS. Specifically, to obtain adaptive prompts, we introduce a multiscale feature extraction strategy to mine rich defect features by the image encoder of SAM. Then, self-correcting probability prototype and cosine similarity attention are proposed to form a mixed similarity aggregation module, which obtains multiscale feature similarities for defect localization. Based on the similarities, a prompt embedding construction module is designed to further extract fine defect location information and generate prompt embeddings for FIDS. Sufficient experiments on MVTec-Unseen, SDD, FSSD-12 and the collected CID datasets demonstrate the superiority of our method.