DiffDD: A surface defect detection framework with diffusion probabilistic model

Abstract

Automated defect detection is critical to improve product quality, reduce costs, and enhance productivity in the manufacturing process. More recently, diffusion-based generative models have entered a new realm of vision tasks by offering high-quality and diverse images. However, the diffusion-based defect detection model fails to be researched, and generative models typically require postprocessing modules to achieve the segmentation of defective regions. Furthermore, generic diffusion models struggle to reconstruct defective regions while preserving the semantic information of the inspection images. To address these issues, we propose a Diffusion-based Defect Detection (DiffDD) framework, comprising a pre-trained backbone (PvTv2) and diffusion probabilistic model (DPM), for surface defect detection. The backbone is employed to extract multi-scale feature maps from input images supervised by pixel-wise labels. These pyramid maps interact with DPM to guide the diffusion process, refining the semantic information of the reconstructed image. Additionally, a semantic restoration (SR) module is used to restore the output of the backbone for loss computation with the pixel-wise label. A conditional image prior method is implemented to constrain the randomness of DPM, aiding the model in generating segmentation masks corresponding to the inspection images. Experimental results on NEU_SEG and MT demonstrate that our model outperforms state-of-the-art (SOTA) methods by a significant margin, indicating the generalization and effectiveness of the proposed DiffDD framework.