Surface Defect Detection with Segmentation-Decision Network on KolektorSDD

Official TensorFlow implementation for Segmentation-based deep-learning approach for surface-defect detection that uses segmentation and decision networks for the detection of surface defects. This work was done in collaboration with Kolektor Group d.o.o..

Citation:

Please cite JIM 2019 journal paper:

@article{Tabernik2019JIM,
  author = {Tabernik, Domen and {\v{S}}ela, Samo and Skvar{\v{c}}, Jure and Sko{\v{c}}aj, Danijel},
  journal = {Journal of Intelligent Manufacturing},
  title = {{Segmentation-Based Deep-Learning Approach for Surface-Defect Detection}},
  year = {2019},
  month = {May},
  day = {15},
  issn={1572-8145},
  doi={10.1007/s10845-019-01476-x}
}

Dependencies:

• python2.7
• TensorFlow r1.1 or newer (tested up to r1.8)
• python libs: numpy, scipy, six, PIL, sklearn, pylab, matplotlib

Dataset

The full dataset Kolektor Surface Defect Dataset (KolektorSDD) is available here.

We split the dataset into three folds to perform 3-fold cross validation. The splits are available at http://box.vicos.si/skokec/gostop/KolektorSDD-training-splits.zip.

Fully prepared TensorFlow dataset split into 3 folds is available at http://box.vicos.si/skokec/gostop/KolektorSDD-dilate=5-tensorflow.zip.

Usage of training/evaluation code

The following files are used to train/evaluate the model:

• segdec_train.py: MAIN ENTRY for training and evaluation
• segdec_model.py: model file for the network
• segdec_data.py: dataset class for training the model

Using the TensorFlow ready KolektorSDD (with dilate=5 for mask) dataset you can train and evaluate with the following:

# 1. Download and extract `KolektorSDD-dilate=5-tensorflow.zip`
mkdir db
cd db
wget http://box.vicos.si/skokec/gostop/KolektorSDD-dilate=5-tensorflow.zip
unzip -x KolektorSDD-dilate=5-tensorflow.zip
cd ..


# Empty folder where models/results will be stored
export OUTPUT_FOLDER=`pwd`/output

# folder where `KolektorSDD-dilate=5-tensorflow.zip` is extracted (must contain `KolektorSDD-dilate=5` subfolder).
export DATASET_FOLDER=`pwd`/db

mkdir $OUTPUT_FOLDER

# 2. Train only segmentation network first:

python -u segdec_train.py --fold=0,1,2 --gpu=0 --max_steps=6600 --train_subset=train \
        --seg_net_type=ENTROPY \
        --size_height=1408 \
        --size_width=512 \
        --with_seg_net=True \
        --with_decision_net=False \
        --storage_dir=$OUTPUT_FOLDER \
        --dataset_dir=$DATASET_FOLDER \
        --datasets=KolektorSDD-dilate=5 \
        --name_prefix=full-size_cross-entropy

# 3. Train and evaluate decision network based on existing segmentation network:

# The `--pretrained_main_folder` must point to the folder where 'fold_XY' subfolders with the trained segmentation models are.
# NOTE: Getting several `Not found: Key tower_0//decision` warrnings when loading the model is OK since the pre-trained model does not have decision net layers yet.

python -u segdec_train.py --fold=0,1,2 --gpu=0 --max_steps=6600 --train_subset=train \
            --seg_net_type=ENTROPY \
            --size_height=1408 \
            --size_width=512 \
            --with_seg_net=False \
            --with_decision_net=True \
            --storage_dir=$OUTPUT_FOLDER \
            --dataset_dir=$DATASET_FOLDER \
            --datasets=KolektorSDD-dilate=5 \
            --name_prefix=decision-net_full-size_cross-entropy \
            --pretrained_main_folder=$OUTPUT_FOLDER/segdec_train/KolektorSDD-dilate=5/full-size_cross-entropy


# 4. Print evaluation metrics combined from all folds

python -u segdec_print_eval.py $OUTPUT_FOLDER/segdec_eval/KolektorSDD-dilate=5/decision-net_full-size_cross-entropy

Note: The model is sensitive to random data shuffles during the training and will lead to different performance with different runs.