基于对比学习与去噪扩散模型的薄膜表面瑕疵图像分类

doi:10.3778/j.issn.1002-8331.2407-0205

摘要/Abstract

摘要： 薄膜材料在生产过程中会出现多种类别的表面瑕疵，薄膜表面瑕疵具有类间差异小、图像数据集不平衡的特点，导致分类性能差。针对上述问题提出了基于对比学习与去噪扩散模型的分类方法。在图像类别数据集上训练去噪扩散模型，得到去噪扩散模型中的噪声预测网络的中间编码输出，即噪声特征；并行地，使用卷积神经网络提取图像特征，并与噪声特征进行融合，得到图像在特征空间中的表达，称为图像的融合特征。使用标签嵌入对比学习方法，将图像类别标签映射到特征空间中作为原型特征，计算融合特征与原型特征之间的对比损失，优化不同类别图像的原型特征在特征空间中的分布，最终得到类别间的特征差异。在锂电池铝塑膜表面瑕疵数据集上进行实验，获得了96.97%的最佳准确率，优于目前的主流方法。

关键词: 对比学习, 度量学习, 去噪扩散模型, 图像分类, 特征融合

Abstract: The surface defects from the manufacturing process of film materials are usually in multiple categories. The properties of small inter-class differences and imbalanced datasets are the main reasons of poor classification performance. A classification method based on contrastive learning and diffusion model is proposed to address the aforementioned issues in this paper. A diffusion model is trained on the image dataset, and the noise features are obtained from the encoded outputs of the noise prediction network which is a part of this diffusion model. On the other hand, the image features are extracted from a CNN model, which are fused with noise features to obtain the enhanced representation of defect images (called the fused feature of image) in the feature space. Using label-embedding contrastive learning to map labels into the feature space, the prototype features are obtained, which are used to calculate the contrastive loss with respect to the fused features in the learning process, and finally the distribution of prototype features of different image categories is optimized in the feature space, which shows the delicate differences among image classes. Experimental validation on classifying surface defects of lithium battery aluminum-plastic films achieves a maximum accuracy of 96.97%, surpassing current mainstream methods.

Key words: contrastive learning, metric learning, diffusion model, image classification, feature fusion

邓皓文, 王恒升. 基于对比学习与去噪扩散模型的薄膜表面瑕疵图像分类[J]. 计算机工程与应用, 2025, 61(21): 242-252.

DENG Haowen, WANG Hengsheng. Image Classification for Film Surface Defect Based on Contrastive Learning and Diffusion Model[J]. Computer Engineering and Applications, 2025, 61(21): 242-252.

参考文献

[1] 李世雄. 基于机器视觉的铝塑膜表面缺陷检测技术研究[D]. 深圳: 深圳大学, 2020.
LI S X. Research on surface defect detection technology of aluminum plastic film based on machine vision[D]. Shenzhen: Shenzhen University, 2020.
[2] 林强, 邬依林, 倪君仪, 等. 基于随机森林的锂电池电极缺陷粗细分类法[J]. 广东技术师范大学学报, 2022, 43(6): 72-77.
LIN Q, WU Y L, NI J Y, et al. A coarse and fine classification method for surface defects of lithium battery electrodes based on random forest[J]. Journal of Guangdong Polytechnic Normal University, 2022, 43(6): 72-77.
[3] 鲁永帅, 唐英杰, 马鑫然, 等. 应用卷积神经网络的锂电池极片涂布缺陷分类[J]. 包装工程, 2022, 43(9): 231-238.
LU Y S, TANG Y J, MA X R, et al. Defect classification of lithium battery pole piece coating using convolutional neural network[J]. Packaging Engineering, 2022, 43(9): 231-238.
[4] CHEON S, LEE H K, KIM C O, et al. Convolutional neural network for wafer surface defect classification and the detection of unknown defect class[J]. IEEE Transactions on Semiconductor Manufacturing, 2019, 32(2): 163-170.
[5] 田志新, 徐震, 茅健, 等. 基于多尺度特征融合的钢材表面缺陷分类方法[J]. 电子科技, 2024, 37(2): 87-94.
TIAN Z X, XU Z, MAO J, et al. Classification method of steel surface defects based on multi-scale feature fusion[J]. Electronic Science and Technology, 2024, 37(2): 87-94.
[6] CHEN T, KORNBLITH S, NOROUZI M, et al. A simple framework for contrastive learning of visual representations[J]. arXiv:2002.05709, 2020.
[7] HE K M, FAN H Q, WU Y X, et al. Momentum contrast for unsupervised visual representation learning[C]//Proceedings of the 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition. Piscataway: IEEE, 2020: 9726-9735.
[8] KHOSLA P, TETERWAK P, WANG C, et al. Supervised contrastive learning[C]//Proceedings of the 34th International Conference on Neural Information Processing Systems, Vancouver, Dec 6-12, 2020. Red Hook, NY, USA: Curran Associates, 2020: 18661-18673.
[9] LIN Z, DU B L, ZHAO X M, et al. Surface defect detection method based on improved semisupervised multitask generative adversarial network[J]. Scientific Programming, 2022, 2022(1): 4481495.
[10] HU J H, YAN P, SU Y T, et al. A method for classification of surface defect on metal workpieces based on twin attention mechanism generative adversarial network[J]. IEEE Sensors Journal, 2021, 21(12): 13430-13441.
[11] AUSTIN J, JOHNSON D D, HO J, et al. Structured denoising diffusion models in discrete state-spaces[C]//Proceedings of the 35th International Conference on Neural Information Processing Systems. Red Hook, NY, USA: Curran Associates, 2024: 17981-17993.
[12] RAMESH A, DHARIWAL P, NICHOL A, et al. Hierarchical text-conditional image generation with CLIP latents[J]. arXiv:2204.06125, 2022.
[13] MUKHOPADHYAY S, GWILLIAM M, AGARWAL V, et al. Diffusion models beat GANs on image classification[J]. arXiv:2307.08702, 2023.
[14] HO J, JAIN A, ABBEEL P. Denoising diffusion probabilistic models[C]//Proceedings of the 34th International Conference on Neural Information Processing Systems, Vancouver, Dec 6-12, 2020. Red Hook, NY, USA: Curran Associates, 2020: 6840-6851.
[15] RONNEBERGER O, FISCHER P, BROX T. U-Net: convolutional networks for biomedical image segmentation[C]// Proceedings of the International Conference on Medical Image Computing and Computer?Assisted Intervention. Cham: Springer, 2015: 234-241.
[16] DHARIWAL P, NICHOL A. Diffusion models beat GANs on image synthesis[C]//Proceedings of the 35th International Conference on Neural Information Processing Systems. Red Hook, NY, USA: Curran Associates, 2024: 8780-8794.
[17] ZADEH A, CHEN M H, PORIA S, et al. Tensor fusion network for multimodal sentiment analysis[C]//Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing. Stroudsburg: ACL, 2017: 1103-1114.
[18] LIU Z, SHEN Y, LAKSHMINARASIMHAN V B, et al. Efficient low-rank multimodal fusion with modality-specific factors[C]//Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics. Stroudsburg: ACL, 2018: 2247-2256.
[19] DOSOVITSKIY A, BEYER L, KOLESNIKOV A, et al. An image is worth 16x16 words: transformers for image recognition at scale[J]. arXiv:2010.11929, 2020.
[20] KORNBLITH S, NOROUZI M, LEE H, et al. Similarity of neural network representations revisited[C]//Proceedings of the 36th International Conference on Machine Learning, Long Beach, Jun 9-15, 2019. San Diego, USA: PMLR, 2019: 3519-3529.