{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/vinhngx/DeepLearningExamples/blob/vinhn_unet_industrial_demo/TensorFlow/Segmentation/UNet_Industrial/notebooks/TensorFlow_UNet_Industrial_TF_train_and_inference.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Gwt7z7qdmTbW"
   },
   "outputs": [],
   "source": [
    "# Copyright 2019 NVIDIA Corporation. All Rights Reserved.\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "# =============================================================================="
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "i4NKCp2VmTbn"
   },
   "source": [
    "<img src=\"http://developer.download.nvidia.com/compute/machine-learning/frameworks/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# UNet Industrial Training and Inference Demo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "fW0OKDzvmTbt"
   },
   "source": [
    "## Overview\n",
    "\n",
    "\n",
    "This U-Net model is adapted from the original version of the [U-Net model](https://arxiv.org/abs/1505.04597) which is\n",
    "a convolutional auto-encoder for 2D image segmentation. U-Net was first introduced by\n",
    "Olaf Ronneberger, Philip Fischer, and Thomas Brox in the paper:\n",
    "[U-Net: Convolutional Networks for Biomedical Image Segmentation](https://arxiv.org/abs/1505.04597).\n",
    "\n",
    "This work proposes a modified version of U-Net, called `TinyUNet` which performs efficiently and with very high accuracy\n",
    "on the industrial anomaly dataset [DAGM2007](https://resources.mpi-inf.mpg.de/conference/dagm/2007/prizes.html).\n",
    "*TinyUNet*, like the original *U-Net* is composed of two parts:\n",
    "- an encoding sub-network (left-side)\n",
    "- a decoding sub-network (right-side).\n",
    "\n",
    "It repeatedly applies 3 downsampling blocks composed of two 2D convolutions followed by a 2D max pooling\n",
    "layer in the encoding sub-network. In the decoding sub-network, 3 upsampling blocks are composed of a upsample2D\n",
    "layer followed by a 2D convolution, a concatenation operation with the residual connection and two 2D convolutions.\n",
    "\n",
    "`TinyUNet` has been introduced to reduce the model capacity which was leading to a high degree of over-fitting on a\n",
    "small dataset like DAGM2007. The complete architecture is presented in the figure below:\n",
    "\n",
    "![UnetModel](https://github.com/vinhngx/DeepLearningExamples/blob/vinhn_unet_industrial_demo/TensorFlow/Segmentation/UNet_Industrial/images/unet.png?raw=1)\n",
    "\n",
    "\n",
    "\n",
    "### Learning objectives\n",
    "\n",
    "This notebook demonstrates the steps for training a UNet model. We then employ the trained model to make inference on new images.\n",
    "\n",
    "## Content\n",
    "1. [Requirements](#1)\n",
    "1. [Data download and preprocessing](#2)\n",
    "1. [Training](#3)\n",
    "1. [Testing trained model](#4)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "aDFrE4eqmTbv"
   },
   "source": [
    "<a id=\"1\"></a>\n",
    "## 1. Requirements\n",
    "\n",
    "\n",
    "### 1.1 Docker container\n",
    "The most convenient way to make use of the NVIDIA Tensorflow UNet model is via a docker container, which provides a self-contained, isolated and re-producible environment for all experiments. Refer to the [Quick Start Guide section](https://github.com/vinhngx/DeepLearningExamples/tree/vinhn_unet_industrial_demo/TensorFlow/Segmentation/UNet_Industrial#requirements) of the Readme documentation for a comprehensive guide. We briefly summarize the steps here.\n",
    "\n",
    "First, clone the repository:\n",
    "\n",
    "```\n",
    "git clone https://github.com/NVIDIA/DeepLearningExamples.git\n",
    "cd DeepLearningExamples/TensorFlow/Segmentation/UNet_Industrial\n",
    "```\n",
    "\n",
    "Next, build the NVIDIA UNet_Industrial container:\n",
    "\n",
    "```\n",
    "docker build . --rm -t unet_industrial:latest\n",
    "```\n",
    "\n",
    "Then launch the container with:\n",
    "\n",
    "```\n",
    "nvidia-docker run -it --rm \\\n",
    "    --shm-size=2g --ulimit memlock=-1 --ulimit stack=67108864 \\\n",
    "    -v /path/to/dataset:/data/dagm2007/ \\\n",
    "    -v /path/to/results:/results \\\n",
    "    unet_industrial:latest\n",
    "```\n",
    "where `/path/to/dataset` is the path on the host machine where the data was/is to be downloaded. More on data set preparation in the next section. `/path/to/results` is wher the trained model will be stored.\n",
    "\n",
    "Within the docker interactive bash session, start Jupyter with\n",
    "\n",
    "```\n",
    "jupyter notebook --ip 0.0.0.0 --port 8888\n",
    "```\n",
    "\n",
    "Then open the Jupyter GUI interface on your host machine at http://localhost:8888. Within the container, this notebook itself is located at `/workspace/unet_industrial/notebooks`.\n",
    "\n",
    "### 1.2 Hardware\n",
    "This notebook can be executed on any CUDA-enabled NVIDIA GPU, although for efficient mixed precision training, a [Tensor Core NVIDIA GPU](https://www.nvidia.com/en-us/data-center/tensorcore/) is desired (Volta, Turing or newer architectures). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "k7RLEcKhmTb0"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mon Sep 30 04:12:27 2019       \r\n",
      "+-----------------------------------------------------------------------------+\r\n",
      "| NVIDIA-SMI 418.40.04    Driver Version: 418.40.04    CUDA Version: 10.1     |\r\n",
      "|-------------------------------+----------------------+----------------------+\r\n",
      "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\r\n",
      "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\r\n",
      "|===============================+======================+======================|\r\n",
      "|   0  Tesla V100-SXM2...  On   | 00000000:86:00.0 Off |                    0 |\r\n",
      "| N/A   41C    P0    60W / 300W |    316MiB / 16130MiB |      0%      Default |\r\n",
      "+-------------------------------+----------------------+----------------------+\r\n",
      "                                                                               \r\n",
      "+-----------------------------------------------------------------------------+\r\n",
      "| Processes:                                                       GPU Memory |\r\n",
      "|  GPU       PID   Type   Process name                             Usage      |\r\n",
      "|=============================================================================|\r\n",
      "+-----------------------------------------------------------------------------+\r\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "HqSUGePjmTb9"
   },
   "source": [
    "<a id=\"2\"></a>\n",
    "## 2. Data download and preprocessing\n",
    "\n",
    "We will first download some data, in particular, the [Weakly Supervised Learning for Industrial Optical Inspection (DAGM 2007)](https://resources.mpi-inf.mpg.de/conference/dagm/2007/prizes.html) dataset. \n",
    "\n",
    "> The competition is inspired by problems from industrial image processing. In order to satisfy their customers' needs, companies have to guarantee the quality of their products, which can often be achieved only by inspection of the finished product. Automatic visual defect detection has the potential to reduce the cost of quality assurance significantly.\n",
    ">\n",
    "> The competitors have to design a stand-alone algorithm which is able to detect miscellaneous defects on various background textures.\n",
    ">\n",
    "> The particular challenge of this contest is that the algorithm must learn, without human intervention, to discern defects automatically from a weakly labeled (i.e., labels are not exact to the pixel level) training set, the exact characteristics of which are unknown at development time. During the competition, the programs have to be trained on new data without any human guidance.\n",
    "\n",
    "**Source:** https://resources.mpi-inf.mpg.de/conference/dagm/2007/prizes.html\n",
    "\n",
    "\n",
    "**Important Information**: The data download script below will download the *public* DAGM 2007 data set, while the *private* DAGM 2007 requires an account to be downloaded. The script will invite you to download them manually and put them in the correct directory for subsequent pre-processing. We will employ the *private* DAGM data set of 10 classes to train UNnet models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "S2PR7weWmTcK"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "################################################\n",
      "Processing Public Dataset\n",
      "################################################\n",
      "\n",
      "Archive:  /data/DeepLearningExamples/TensorFlow/Segmentation/UNet_Industrial/notebooks/data/zip_files/public/Class1.zip\n",
      "Archive:  /data/DeepLearningExamples/TensorFlow/Segmentation/UNet_Industrial/notebooks/data/zip_files/public/Class1_def.zip\n",
      "Archive:  /data/DeepLearningExamples/TensorFlow/Segmentation/UNet_Industrial/notebooks/data/zip_files/public/Class2.zip\n",
      "Archive:  /data/DeepLearningExamples/TensorFlow/Segmentation/UNet_Industrial/notebooks/data/zip_files/public/Class2_def.zip\n",
      "Archive:  /data/DeepLearningExamples/TensorFlow/Segmentation/UNet_Industrial/notebooks/data/zip_files/public/Class3.zip\n",
      "Archive:  /data/DeepLearningExamples/TensorFlow/Segmentation/UNet_Industrial/notebooks/data/zip_files/public/Class3_def.zip\n",
      "Archive:  /data/DeepLearningExamples/TensorFlow/Segmentation/UNet_Industrial/notebooks/data/zip_files/public/Class4.zip\n",
      "Archive:  /data/DeepLearningExamples/TensorFlow/Segmentation/UNet_Industrial/notebooks/data/zip_files/public/Class4_def.zip\n",
      "Archive:  /data/DeepLearningExamples/TensorFlow/Segmentation/UNet_Industrial/notebooks/data/zip_files/public/Class5.zip\n",
      "Archive:  /data/DeepLearningExamples/TensorFlow/Segmentation/UNet_Industrial/notebooks/data/zip_files/public/Class5_def.zip\n",
      "Archive:  /data/DeepLearningExamples/TensorFlow/Segmentation/UNet_Industrial/notebooks/data/zip_files/public/Class6.zip\n",
      "Archive:  /data/DeepLearningExamples/TensorFlow/Segmentation/UNet_Industrial/notebooks/data/zip_files/public/Class6_def.zip\n",
      "\n",
      "################################################\n",
      "Processing Private Dataset\n",
      "################################################\n",
      "\n",
      "\n",
      "Error! File not found: '/data/DeepLearningExamples/TensorFlow/Segmentation/UNet_Industrial/notebooks/data/zip_files/private/Class1.zip'\n",
      "Instructions:\n",
      "    1. Create an account on 'https://hci.iwr.uni-heidelberg.de/node/3616'\n",
      "    2. Download the missing file(s) to: '/data/DeepLearningExamples/TensorFlow/Segmentation/UNet_Industrial/notebooks/data/zip_files/private'\n",
      "\n"
     ]
    }
   ],
   "source": [
    "! ../download_and_preprocess_dagm2007.sh ./data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "EQAIszkxmTcT"
   },
   "source": [
    "Within the docker container, the final data directory should look like:\n",
    "\n",
    "```\n",
    "./data\n",
    "  raw_images\n",
    "      public\n",
    "          Class1\t    \n",
    "          Class2\t\n",
    "          Class3\t    \n",
    "          Class4\t\n",
    "          Class5\t    \n",
    "          Class6\n",
    "          Class1_def  \n",
    "          Class2_def\t\n",
    "          Class3_def  \n",
    "          Class4_def\t\n",
    "          Class5_def  \n",
    "          Class6_def\n",
    "      private\n",
    "          Class1  \n",
    "          Class10  \n",
    "          Class2  \n",
    "          Class3  \n",
    "          Class4  \n",
    "          Class5  \n",
    "          Class6  \n",
    "          Class7  \n",
    "          Class8  \n",
    "          Class9\n",
    "  zip_files\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "RL8d9IwzmTcV"
   },
   "source": [
    "<a id=\"3\"></a>\n",
    "## 3. Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "o6wayGf1mTcX"
   },
   "source": [
    "The repository provides several training recipes with 1, 4 and 8 GPU with FP32 and automatic mixed precision in `./script`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "eQPsExf6mTca"
   },
   "source": [
    "### 3.1 Training with 1 GPU\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "HapDsY4VmTce"
   },
   "source": [
    "#### Training with full precision\n",
    "Training on 1 GPU with FP32 with the following syntax:\n",
    "\n",
    "```\n",
    "./UNet_FP32_1GPU.sh <path to result repository> <path to dataset> <DAGM2007 classID (1-10)\n",
    "```\n",
    "\n",
    "For example: \n",
    "```\n",
    "../scripts/UNet_FP32_1GPU.sh ./results/1GPU-FP32 /CDR/DAGM10 1\n",
    "```\n",
    "                        \n",
    "Note that inside the shell script it make use of the main python trains script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RESULT_DIR=\"./results/1GPU-FP32\"\n",
    "DATA_DIR=\"/CDR/DAGM10\"\n",
    "CLASS=\"1\"\n",
    "%run ../main.py \\\n",
    "    --unet_variant='tinyUNet' \\\n",
    "    --activation_fn='relu' \\\n",
    "    --exec_mode='train_and_evaluate' \\\n",
    "    --iter_unit='batch' \\\n",
    "    --num_iter=2500 \\\n",
    "    --batch_size=16 \\\n",
    "    --warmup_step=10 \\\n",
    "    --results_dir={RESULT_DIR} \\\n",
    "    --data_dir={DATA_DIR} \\\n",
    "    --dataset_name='DAGM2007' \\\n",
    "    --dataset_classID={CLASS} \\\n",
    "    --data_format='NCHW' \\\n",
    "    --use_auto_loss_scaling \\\n",
    "    --noamp \\\n",
    "    --noxla \\\n",
    "    --learning_rate=1e-4 \\\n",
    "    --learning_rate_decay_factor=0.8 \\\n",
    "    --learning_rate_decay_steps=500 \\\n",
    "    --rmsprop_decay=0.9 \\\n",
    "    --rmsprop_momentum=0.8 \\\n",
    "    --loss_fn_name='adaptive_loss' \\\n",
    "    --weight_decay=1e-5 \\\n",
    "    --weight_init_method='he_uniform' \\\n",
    "    --augment_data \\\n",
    "    --display_every=250 \\\n",
    "    --debug_verbosity=0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "j-aFEwb4mTcn"
   },
   "source": [
    "#### Training with mixed-precision\n",
    "We now launch the training process using mixed precision. Training on 1 GPU with automatic mixed precision (AMP) with the following syntax:\n",
    "\n",
    "```\n",
    "./UNet_AMP_1GPU.sh <path to result repository> <path to dataset> <DAGM2007 classID (1-10)\n",
    "```\n",
    "\n",
    "For example: \n",
    "```\n",
    "../scripts/UNet_AMP_1GPU.sh ./results/1GPU-AMP/CDR/DAGM10 1\n",
    "```\n",
    "                        \n",
    "Note that inside the shell script it make use of the main python trains script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "o3AZ-CXYmTcp",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "RESULT_DIR=\"./results/1GPU-FP16\"\n",
    "DATA_DIR=\"/CDR/DAGM10\"\n",
    "CLASS=\"1\"\n",
    "%run ../main.py \\\n",
    "    --unet_variant='tinyUNet' \\\n",
    "    --activation_fn='relu' \\\n",
    "    --exec_mode='train_and_evaluate' \\\n",
    "    --iter_unit='batch' \\\n",
    "    --num_iter=2500 \\\n",
    "    --batch_size=16 \\\n",
    "    --warmup_step=10 \\\n",
    "    --results_dir={RESULT_DIR} \\\n",
    "    --data_dir={DATA_DIR} \\\n",
    "    --dataset_name='DAGM2007' \\\n",
    "    --dataset_classID={CLASS} \\\n",
    "    --data_format='NCHW' \\\n",
    "    --use_auto_loss_scaling \\\n",
    "    --amp \\\n",
    "    --noxla \\\n",
    "    --learning_rate=1e-4 \\\n",
    "    --learning_rate_decay_factor=0.8 \\\n",
    "    --learning_rate_decay_steps=500 \\\n",
    "    --rmsprop_decay=0.9 \\\n",
    "    --rmsprop_momentum=0.8 \\\n",
    "    --loss_fn_name='adaptive_loss' \\\n",
    "    --weight_decay=1e-5 \\\n",
    "    --weight_init_method='he_uniform' \\\n",
    "    --augment_data \\\n",
    "    --display_every=250 \\\n",
    "    --debug_verbosity=0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2. Training with 8 GPUs\n",
    "\n",
    "We provide training recipes for 8 GPUs using FP32 or automatic mixed precision (AMP)\n",
    "\n",
    "#### Training with full precision\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!../scripts/UNet_FP32_8GPU.sh ./results/8GPU-FP32 /CDR/DAGM10 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Training with automatic mixed precision (AMP)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!../scripts/UNet_AMP_8GPU.sh ./results/8GPU-AMP /CDR/DAGM10 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "X959LYwjmTcw"
   },
   "source": [
    "<a id=\"4\"></a>\n",
    "## 4. Testing trained model\n",
    "\n",
    "After model training has completed, we can test the trained model against the DAGM 2007 public data set which wasn't used for training. First, we load some required libraries and define some helper functions to load images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TURMzc6HmTcy",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.insert(0,'..')\n",
    "\n",
    "from model.unet import UNet_v1\n",
    "\n",
    "import numpy as np\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.image as mpimg\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f26c924df28>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 661644 characters]\n",
      "text/plain": [
       "<Figure size 720x720 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "img = mpimg.imread('./data/raw_images/public/Class1_def/1.png')\n",
    "\n",
    "plt.figure(figsize = (10,10))\n",
    "plt.imshow(img, cmap='gray')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see in this figure, there exists a defective area in the top left corner. We will now load the model and carry out inference on the normalized test image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Image preprocessing\n",
    "img =  np.expand_dims(img, axis=2)\n",
    "img =  np.expand_dims(img, axis=0)\n",
    "img = (img-0.5)/0.5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we start a TF session, load the trained UNet model and carry out inference on the test image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "checkpoint\r\n",
      "eval\r\n",
      "events.out.tfevents.1569830054.ac4fd7ffba26\r\n",
      "graph.pbtxt\r\n",
      "model.ckpt-0.data-00000-of-00002\r\n",
      "model.ckpt-0.data-00001-of-00002\r\n",
      "model.ckpt-0.index\r\n",
      "model.ckpt-0.meta\r\n",
      "model.ckpt-1000.data-00000-of-00002\r\n",
      "model.ckpt-1000.data-00001-of-00002\r\n",
      "model.ckpt-1000.index\r\n",
      "model.ckpt-1000.meta\r\n",
      "model.ckpt-2000.data-00000-of-00002\r\n",
      "model.ckpt-2000.data-00001-of-00002\r\n",
      "model.ckpt-2000.index\r\n",
      "model.ckpt-2000.meta\r\n",
      "model.ckpt-2500.data-00000-of-00002\r\n",
      "model.ckpt-2500.data-00001-of-00002\r\n",
      "model.ckpt-2500.index\r\n",
      "model.ckpt-2500.meta\r\n"
     ]
    }
   ],
   "source": [
    "!ls ./results/1GPU-FP16/checkpoints  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "config = tf.ConfigProto()\n",
    "config.gpu_options.allow_growth = True\n",
    "config.allow_soft_placement = True\n",
    "\n",
    "graph = tf.Graph()\n",
    "with graph.as_default():\n",
    "    with tf.Session(config=config) as sess:\n",
    "        network = UNet_v1(\n",
    "            model_name=\"UNet_v1\",\n",
    "            input_format='NHWC',\n",
    "            compute_format='NHWC',\n",
    "            n_output_channels=1,\n",
    "            unet_variant='tinyUNet',\n",
    "            weight_init_method='he_uniform',\n",
    "            activation_fn='relu'\n",
    "        )\n",
    "        \n",
    "        tf_input = tf.placeholder(tf.float32, [None, 512, 512, 1], name='input')\n",
    "        \n",
    "        outputs, logits = network.build_model(tf_input)\n",
    "        saver = tf.train.Saver()\n",
    "\n",
    "        # Restore variables from disk.\n",
    "        saver.restore(sess, \"./results/1GPU-FP16/checkpoints/model.ckpt-2500\")\n",
    "        \n",
    "        \n",
    "        output = sess.run([outputs, logits], feed_dict={tf_input: img})\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f2720273be0>"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 10420 characters]\n",
      "text/plain": [
       "<Figure size 720x720 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Print out model predicted mask\n",
    "plt.figure(figsize = (10,10))\n",
    "plt.imshow(np.squeeze(output[0]), cmap='gray')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, the model points out the correct defective area in this image. Please feel free to try out other defective images within `./data/raw_images/public/Class1_def/`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.png\t 114.png  13.png   145.png  25.png  40.png  56.png  71.png  87.png\r\n",
      "10.png\t 115.png  130.png  146.png  26.png  41.png  57.png  72.png  88.png\r\n",
      "100.png  116.png  131.png  147.png  27.png  42.png  58.png  73.png  89.png\r\n",
      "101.png  117.png  132.png  148.png  28.png  43.png  59.png  74.png  9.png\r\n",
      "102.png  118.png  133.png  149.png  29.png  44.png  6.png   75.png  90.png\r\n",
      "103.png  119.png  134.png  15.png   3.png   45.png  60.png  76.png  91.png\r\n",
      "104.png  12.png   135.png  150.png  30.png  46.png  61.png  77.png  92.png\r\n",
      "105.png  120.png  136.png  16.png   31.png  47.png  62.png  78.png  93.png\r\n",
      "106.png  121.png  137.png  17.png   32.png  48.png  63.png  79.png  94.png\r\n",
      "107.png  122.png  138.png  18.png   33.png  49.png  64.png  8.png   95.png\r\n",
      "108.png  123.png  139.png  19.png   34.png  5.png   65.png  80.png  96.png\r\n",
      "109.png  124.png  14.png   2.png    35.png  50.png  66.png  81.png  97.png\r\n",
      "11.png\t 125.png  140.png  20.png   36.png  51.png  67.png  82.png  98.png\r\n",
      "110.png  126.png  141.png  21.png   37.png  52.png  68.png  83.png  99.png\r\n",
      "111.png  127.png  142.png  22.png   38.png  53.png  69.png  84.png  labels.txt\r\n",
      "112.png  128.png  143.png  23.png   39.png  54.png  7.png   85.png\r\n",
      "113.png  129.png  144.png  24.png   4.png   55.png  70.png  86.png\r\n"
     ]
    }
   ],
   "source": [
    "!ls ./data/raw_images/public/Class1_def/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "g8MxXY5GmTc8"
   },
   "source": [
    "# Conclusion\n",
    "\n",
    "In this notebook, we have walked through the complete process of preparing the container and data required for training the UNet-Industrial models. We have also investigated various training options with FP32 and automatic mixed precision, trained and tested UNet models with new test images.\n",
    "\n",
    "## What's next\n",
    "Now it's time to try the UNet model on your own data. Observe the performance impact of mixed precision training while comparing the final accuracy of the models trained with FP32 and mixed precision.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "249yGNLmmTc_"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "include_colab_link": true,
   "name": "TensorFlow_UNet_Industrial_Colab_train_and_inference.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
