Uncertainty-aware and dynamically-mixed pseudo-labels for semi-supervised defect segmentation

Abstract

Deep learning-based defect segmentation is one of the important tasks of machine vision in automated inspection. Supervised learning methods have recently achieved remarkable performance on this task. However, the effectiveness of the supervised methods is limited by the scarcity and high cost of pixel-level annotation of training data. Semi-supervised learning methods have been proposed for training deep learning networks using a limited amount of labeled data along with additional unlabeled data for image segmentation. Most of these methods are based on consistency regularization and pseudo labeling, where the predictions on unlabeled samples often come with noise and are unreliable, resulting in poor segmentation performance. To alleviate this problem, we propose uncertainty-aware pseudo labels, which are generated from dynamically mixed predictions of multiple decoders that leverage a shared encoder network. The estimated uncertainty guides the pseudo-label-based supervision and regularizes the training when using the unlabeled samples. In our experiments on four public datasets for defect segmentation, the proposed method outperformed the fully supervised baseline and six state-of-the-art semi-supervised segmentation methods. We also conducted an extensive ablation study to demonstrate the effectiveness of our approach in various settings. The implementation code for this work is available at https://github.com/djene-mengistu/UAPS.