djene-mengistu / UAPS

This repo contains implementation of uncertainty estimation, rectification, and minimization for guiding the pseudo-label learning in semi-supervised defect segmentation setting.

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.

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.

Full paper source:

You can read the details about the methods, implementation, and results from: (https://www.sciencedirect.com/science/article/pii/S0166361523001458)

Please cite ourwork 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},
}

Python >= 3.6

PyTorch >= 1.1.0 PyYAML, tqdm, tensorboardX

Data Preparation

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

Code usage

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.py

To test the performance of the proposed method, load the trained weights accordingly and run the following:

run UAPS_Testing.ipynb

Similarly, train the proposed method, UAPS, for the other datasets from the indicated directories after setting appropriate hyper-parametres.

Some results and visualizations

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.

Ablation experiments

Results from different model settings are presented as follows:

Fig. 6: Effects of different loss combinations.

Fig. 7: Effects of loss coefficents.

Fig. 8: Proposed dynamic-mixing vs. Averaging to generate pseudo-labels.

Fig. 9: Effects of number of auxuliary decoders on segmentation performance and inference time.

Acknowledgment

This repo borrowed many implementations from SSL4MIS and simEps

Contact

For any issue please contact me at [e-mail]