{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Inference_visualization_tool.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Klhdy8pnk5J8"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/tensorflow/models/blob/master/official/projects/edgetpu/vision/serving/inference_visualization_tool.ipynb)\n",
        "\n",
        "# Visualizing segmentation outputs using colab.\n",
        "\n",
        "This file is located in [github](https://github.com/tensorflow/models/blob/master/official/projects/edgetpu/vision/serving/inference_visualization_tool.ipynb) and uses [colab integration](https://colab.research.google.com/github/tensorflow/models/blob/master/official/projects/edgetpu/vision/serving/inference_visualization_tool.ipynb) to seemlessly show [segmentation model](https://github.com/tensorflow/models/blob/master/official/projects/edgetpu/vision/README.md) outputs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dHWaHWjYJEsE"
      },
      "source": [
        "## Setup sandbox\n",
        "\n",
        "Imports required libs and get ready to load data."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YWVHsuCiJHjo"
      },
      "source": [
        "from google.colab import auth # access to saved model in tflite format\n",
        "auth.authenticate_user()\n",
        "from PIL import Image # used to read images as arrays\n",
        "import tensorflow as tf # runs tested model\n",
        "import numpy as np # postprocessing for render.\n",
        "from scipy import ndimage # postprocessing for render.\n",
        "import matplotlib.pyplot as plt # render\n",
        "\n",
        "# Copies reference to colab's sandbox.\n",
        "def copy_to_sandbox(web_path):\n",
        "  sandbox_path = web_path.split('/')[-1]\n",
        "  !rm -f {sandbox_path}\n",
        "  if web_path[:2] == \"gs\":\n",
        "    !gsutil cp {web_path} {sandbox_path}\n",
        "  else:\n",
        "    !wget -v {web_path} --no-check-certificate\n",
        "  return sandbox_path\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BzsTQjiEI75t"
      },
      "source": [
        "## Prepare sandbox images\n",
        "\n",
        "Running this notebook will show sample segmentation of 3 pictures from [ADE20K](https://groups.csail.mit.edu/vision/datasets/ADE20K/) dataset. You can try it on other pictures by adding you own URLS to `IMAGE_URLS` list."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-vGHZSPWXbyu"
      },
      "source": [
        "# Common image URL pattern.\n",
        "_IMAGE_URL_PATTERN = 'https://raw.githubusercontent.com/tensorflow/models/master/official/projects/edgetpu/vision/serving/testdata/ADE_val_{name}.jpg'\n",
        "# Coma separated list of image ids.\n",
        "_IMAGE_NAMES = ['00001626','00001471','00000557']\n",
        "# List\n",
        "IMAGE_URLS = [_IMAGE_URL_PATTERN.replace('{name}', image) for image in _IMAGE_NAMES]\n",
        "# IMAGE_URLS.append('your URL')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HMgtac7ULZcK"
      },
      "source": [
        "IMAGES = [copy_to_sandbox(image_url) for image_url in IMAGE_URLS]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YJU3mFebL6oP"
      },
      "source": [
        "## Prepare sandbox model\n",
        "\n",
        "Default visualize is running M-size model. Model is copiend to sandbox to run.You can use another model from the list."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PSQPTYgLL_VS"
      },
      "source": [
        "MODEL_HOME='gs://tf_model_garden/models/edgetpu/checkpoint_and_tflite/vision/segmentation-edgetpu/tflite/default_argmax'\n",
        "!gsutil ls {MODEL_HOME}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "shzpQoEaGnvp"
      },
      "source": [
        "# Path to tflite file, can use any other from list above.\n",
        "MODEL_NAME='deeplabv3plus_mobilenet_edgetpuv2_m_ade20k_32.tflite'#@param\n",
        "MODEL = copy_to_sandbox(MODEL_HOME + \"/\" + MODEL_NAME)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6Bd4if7fMY7v"
      },
      "source": [
        "# Image sizes compatible with the model\n",
        "MODEL_IMAGE_WIDTH = 512\n",
        "MODEL_IMAGE_HEIGHT = 512"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mtV4GhuXQn2Q"
      },
      "source": [
        "## Image preprocess\n",
        "\n",
        "Function defines how to preprocess image before running inference"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jASE46vxRHeP"
      },
      "source": [
        "def read_image(image):\n",
        "  im = Image.open(image).convert('RGB')\n",
        "  min_dim=min(im.size[0], im.size[1])\n",
        "  new_y_dim = MODEL_IMAGE_HEIGHT * im.size[0] // min_dim\n",
        "  new_x_dim = MODEL_IMAGE_WIDTH * im.size[1] // min_dim\n",
        "  # scale to outer fit.\n",
        "  im = im.resize((new_y_dim, new_x_dim))\n",
        "  input_data = np.expand_dims(im, axis=0)\n",
        "  # crop to size\n",
        "  return input_data[:, :MODEL_IMAGE_HEIGHT, :MODEL_IMAGE_WIDTH]\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Kkxj-SkrNZE2"
      },
      "source": [
        "## Model runner.\n",
        "\n",
        "Simple wrapper of tflite interpreter  invoke."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GdlsbiVqL5JZ"
      },
      "source": [
        "def run_model(input_data, model_data):\n",
        "  preprocessed_data = (input_data-128).astype(np.int8)\n",
        "  # Load the tflite model and allocate tensors.\n",
        "  interpreter_x = tf.lite.Interpreter(model_path=model_data)\n",
        "  interpreter_x.allocate_tensors()\n",
        "  # Get input and output tensors.\n",
        "  input_details = interpreter_x.get_input_details()\n",
        "  output_details = interpreter_x.get_output_details()\n",
        "  interpreter_x.set_tensor(input_details[0]['index'], preprocessed_data)\n",
        "  interpreter_x.invoke()\n",
        "  output_data = interpreter_x.get_tensor(output_details[0]['index'])\n",
        "  return output_data.reshape((MODEL_IMAGE_HEIGHT, MODEL_IMAGE_WIDTH))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KrJYUoRYOShc"
      },
      "source": [
        "## 6px wide edge highlighter.\n",
        "\n",
        "First function bellow finds edges of classes, and highlights them with 6px edge. Second function blends edge with original image."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KhS1lOrxHp5C"
      },
      "source": [
        "# Creates a 6px wide boolean edge mask to highlight the segmentation.\n",
        "def edge(mydata):\n",
        "  mydata = mydata.reshape((MODEL_IMAGE_HEIGHT, MODEL_IMAGE_WIDTH))\n",
        "  mydatat = mydata.transpose([1, 0])\n",
        "  mydata = np.convolve(mydata.reshape(-1), [-1, 0, 1], mode='same').reshape((MODEL_IMAGE_HEIGHT, MODEL_IMAGE_WIDTH))\n",
        "  mydatat = np.convolve(mydatat.reshape(-1), [-1, 0, 1], mode='same').reshape((MODEL_IMAGE_HEIGHT, MODEL_IMAGE_WIDTH)).transpose([1, 0])\n",
        "  mydata = np.maximum((mydata != 0).astype(np.int8), (mydatat != 0).astype(np.int8))\n",
        "  mydata = ndimage.binary_dilation(mydata).astype(np.int8)\n",
        "  mydata = ndimage.binary_dilation(mydata).astype(np.int8)\n",
        "  mydata = ndimage.binary_dilation(mydata).astype(np.int8)\n",
        "  return mydata"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GRxyl3DkSeIF"
      },
      "source": [
        "def fancy_edge_overlay(input_data, output_data):\n",
        "  output_data = np.reshape(np.minimum(output_data, 32), (MODEL_IMAGE_HEIGHT, MODEL_IMAGE_WIDTH))\n",
        "  output_edge = edge(output_data).reshape((MODEL_IMAGE_HEIGHT, MODEL_IMAGE_WIDTH,1))\n",
        "  output_data = np.stack([output_data%3, (output_data//3)%3, (output_data//9)%3], axis = -1)\n",
        "  return input_data.reshape((MODEL_IMAGE_HEIGHT, MODEL_IMAGE_WIDTH, 3)).astype(np.float32) * (1-output_edge) + output_data * output_edge * 255\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1NJY0tKaTjdW"
      },
      "source": [
        "## Visualize!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1mot5M_nl5P7"
      },
      "source": [
        "# Set visualization wind sizes.\n",
        "fig, ax = plt.subplots(max(len(IMAGES),2), 3)\n",
        "fig.set_figwidth(30)\n",
        "fig.set_figheight(10*max(len(IMAGES),2))\n",
        "\n",
        "# Read and test image.\n",
        "for r, image in enumerate(IMAGES):\n",
        "  input_data = read_image(image)\n",
        "  ax[r, 0].imshow(input_data.reshape((MODEL_IMAGE_HEIGHT, MODEL_IMAGE_WIDTH, 3)).astype(np.uint8))\n",
        "  ax[r, 0].set_title('Original')\n",
        "  ax[r, 0].grid(False)\n",
        "\n",
        "  # Test the model on input data.\n",
        "  output_data = run_model(input_data, MODEL)\n",
        "  ax[r, 1].imshow(output_data, vmin = 0, vmax = 32)\n",
        "  ax[r, 1].set_title('Segmentation')\n",
        "  ax[r, 1].grid(False)\n",
        "\n",
        "  fancy_data = fancy_edge_overlay(input_data, output_data)\n",
        "  ax[r, 2].imshow(fancy_data.astype(np.uint8), vmin = 0, vmax = 32)\n",
        "  ax[r, 2].set_title('Segmentation & original')\n",
        "  ax[r, 2].grid(False)\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}