Uncertainty-aware and Dynamically-mixed Pseudo-labels for Semi-supervised Defect Segmentation (UAPS)
In this study, we propose a novel 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.
You can read the details about the methods, implementation, and results from: (https://www.sciencedirect.com/science/article/pii/S0166361523001458)
🔥 If you find it useful, also cite our work as follows:
@article{SIME2023103995,
title = {Uncertainty-aware and dynamically-mixed pseudo-labels for semi-supervised defect segmentation},
journal = {Computers in Industry},
volume = {152},
pages = {103995},
year = {2023},
issn = {0166-3615},
doi = {https://doi.org/10.1016/j.compind.2023.103995},
url = {https://www.sciencedirect.com/science/article/pii/S0166361523001458},
author = {Dejene M. Sime and Guotai Wang and Zhi Zeng and Bei Peng},
}
The illustration of the uncertainty map and error is presented as follows:
Fig. 1: Illustration of the network prediction, prediction error which is false negatives and false positives, and the uncertainty map computed as the KL-distance between primary decoder and average predictions.
The overall architecture of the proposed method is presented as follows:
Fig. 2: Illustration of the overall framework of the UAPS using the unlabeled samples. We apply various perturbations to the encoder output to obtain different predictions with multiple decoders. The predictions from the decoders are then combined using randomly generated weights to compute a pseudo-label. We calculate uncertainty as the KL-distance between the average prediction and each decoder’s prediction.
The baseline architecture and the perterbation method is presented as follows:
Fig. 3: Illustration of baseline architecture adapted from U-Net with perturbation module. The encoder output remains unchanged at each block, while the type of perturbation changes to produce different versions of the outputs to be fed to the decoder networks.
Download datasets. There are 4 datasets to download:
- NEU-SEG dataset from NEU-seg
- DAGM dataset from DAGM
- MT (Magnetic Tiles) dataset from MTiles
- KoSDD2 (KolektorSDD2) dataset from KoSDD2
Put downloaded data into the following directory structure:
- .../data/
- NEU_data/ ... # raw data of NEU-Seg
- DAGM_data/ ...# raw data of DAGM
- MTiles_data/ ...# raw data of MTiles
- KoSDD2_data/ ...# raw data of KoSDD2
The training files and settings for each compared network is presented in separate directory. Train each network and test from the presented directory. To train the proposed UAPS method run the following after setting hyperparameters such as labeled-ratio, iteration-per-epoch, consistency ramp length, and consistency loss coefficients.
python UAPS_train.pyTo test the performance of the proposed method, load the trained weights accordingly and run the following:
run UAPS_Testing.ipynbSimilarly, train the proposed method, UAPS, for the other datasets from the indicated directories after setting appropriate hyper-parametres.
Visualization of the segmetnation results on selected samples are presented as follows:
Fig. 4: Visualization of the segmentation results on the NEU-Seg dataset. The regions indicated by the dashed-red-box shows wrong prediction.
Fig. 5: Visualization of the segmentation results on the MTiles dataset.
This repo borrowed many implementations from simEps
For any issue please contact me at djene.mengistu@gmail.com