UNet_Industrial TensorFlow checkpoint trained with AMP
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
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.