Model Overview

This model is a convolutional neural network for 2D image segmentation tuned to avoid overfitting.

Model Architecture

This UNet model is adapted from the original version of the UNet model which is a convolutional auto-encoder for 2D image segmentation. UNet was first introduced by Olaf Ronneberger, Philip Fischer, and Thomas Brox in the paper: UNet: Convolutional Networks for Biomedical Image Segmentation.

This work proposes a modified version of UNet, called TinyUNet which performs efficiently and with very high accuracy on the industrial anomaly dataset DAGM2007. TinyUNet, like the original UNet is composed of two parts:

• an encoding sub-network (left-side)
• a decoding sub-network (right-side).

It repeatedly applies 3 downsampling blocks composed of two 2D convolutions followed by a 2D max pooling layer in the encoding sub-network. In the decoding sub-network, 3 upsampling blocks are composed of a upsample2D layer followed by a 2D convolution, a concatenation operation with the residual connection and two 2D convolutions.

TinyUNet has been introduced to reduce the model capacity which was leading to a high degree of over-fitting on a small dataset like DAGM2007. The complete architecture is presented in the figure below:

Figure 1. Architecture of the UNet Industrial

Training

This model was trained using script available on NGC and in GitHub repo

Dataset

The following datasets were used to train this model:

• DAGM2007 - Synthetic dataset for defect detection on textured surfaces.

Performance

Performance numbers for this model are available in NGC

References

• original paper
• NVIDIA model implementation in NGC
• NVIDIA model implementation on GitHub

License

This model was trained using open-source software available in Deep Learning Examples repository. For terms of use, please refer to the license of the script and the datasets the model was derived from.