MetalDefect-SD: Adapting Stable Diffusion for Few-Shot Metal Surface Defect Automatic Segmentation

Abstract

Few-shot segmentation of metal surface defects is challenging because pixel-level annotations are scarce, defect appearances vary widely, and reflective backgrounds often contain defect-like textures. Existing support-query methods and diffusion-based segmentation frameworks can exploit pretrained visual priors, but their support transfer may propagate unreliable background cues into the query representation. We propose MetalDefect-SD, a Stable Diffusion-based framework for support-conditioned defect mask prediction. Its main contribution is the Reliability-Guided Support-Query Interaction (RG-SQI) module, which combines a Defect-Background Self-Support Prototype (DBSSP) branch with Masked Gated Support-to-Query (MG-S2Q) Attention. DBSSP constructs complementary defect and background prototypes from the query latent to improve defect-background discrimination, while MG-S2Q uses the support mask as an attention bias and a learnable gate to regulate support feature propagation. We further adapt a Visual-to-Textual Condition Translator (VTCT) to map annotated support-region features into implicit text-compatible embeddings, enabling the diffusion U-Net to exploit conditional cross-attention without handcrafted prompts. The adapted U-Net predicts the clean query-mask latent in a single step, followed by decoding with a frozen variational autoencoder (VAE). On FSSD-12, MetalDefect-SD achieves 58.5% mean intersection-over-union (mIoU) and 72.5% foreground-background intersection-over-union (FB-IoU) in the 1-shot setting, and 59.0% mIoU and 74.8% FB-IoU in the 5-shot setting. These results demonstrate that reliability-aware support-query interaction improves foreground-background discrimination and enables effective diffusion-prior adaptation for data-efficient industrial defect segmentation. More broadly, this work demonstrates that pretrained diffusion priors can be adapted to data-efficient pixel-level industrial inspection when defect annotations are scarce. The framework therefore connects foundation-model adaptation, few-shot dense prediction, and automated manufacturing quality control.