onnx_to_tflite.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "onnx-to-tflite.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "uZ-iKIlAGFgq",
        "wC6AVy57GP0x",
        "7_1Pxu96XKpU",
        "qx74mzYZFew6",
        "-s-8GlTjiErh",
        "C3RjanvwFt7I",
        "zJ3h-XayF68_",
        "Qn1HMaJGdZnb",
        "DbzMl8aNdaIk"
      ],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/joshuacwnewton/ONNX-to-TFLite/blob/master/onnx_to_tflite.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0HTlBXvCAEi8",
        "colab_type": "text"
      },
      "source": [
        "# 1 Introduction\n",
        "\n",
        "This notebook demonstrates the conversion process from an **.ONNX** model _(exported from MATLAB)_ to a **.tflite** model _(to be used within TensorFlow Lite, on an Android or iOS device.)_ In addition to conversion, this notebook contains cells for running inference using a set of test images to validate that predictions remain consistent across converted models.\n",
        "\n",
        "> **Note:** TensorFlow's API is constantly evolving. This notebook was written in November of 2019, during the transition period from TF 1.X to TF 2.X, so it is likely that relevant APIs will have updated since."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y4h6nfXPxs7g",
        "colab_type": "text"
      },
      "source": [
        "### 1.1 Initializing input files\n",
        "\n",
        "Several files are used to demonstrate the conversion process.\n",
        "\n",
        "- `hasCircularShape_chartObjects_googlenet.onnx`, a demo model exported from MATLAB.\n",
        "- `test_images.zip`, a .zip with two class directories, each containing 25 images.\n",
        "\n",
        "The files themselves are provided in the associated GitHub repository for this notebook.\n",
        "\n",
        "Once uploaded, please validate that the files have been stored in the notebook correctly by running the cell below. Cells can be run by first clicking on them, then using either the **\"Run Cell\"** button ( ▷ ), or typing **Ctrl**+**Enter**.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A71kIPZgAbFo",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Set path variables\n",
        "onnx_path = 'hasCircularShape_chartObjects_googlenet.onnx'\n",
        "img_zip_path = 'test_images.zip'\n",
        "\n",
        "# Check that correct files have been uploaded\n",
        "import os\n",
        "\n",
        "assert os.path.exists(onnx_path)\n",
        "assert os.path.exists(img_zip_path) \n",
        "\n",
        "print(\"Files uploaded successfully.\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "plTTzpBd_-95",
        "colab_type": "text"
      },
      "source": [
        "# 2 ONNX (.onnx) -> TensorFlow FrozenGraph (.pb)\n",
        "\n",
        "Now that the .onnx model file is stored within the notebook, it can be converted to a .pb model file for use within TensorFlow."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "brSS6PxZ__bV",
        "colab_type": "text"
      },
      "source": [
        "### 2.1 Background Information\n",
        "\n",
        "**ONNX** is an open-source format for AI models created Facebook and Microsoft [[1]](https://onnx.ai/). The goal of the ONNX format is to provide interoperability between frameworks. The ONNX project provides conversion tools between the ONNX format and formats from other frameworks [[2]](http://onnx.ai/supported-tools). \n",
        "\n",
        "**MATLAB** allows model exporting to a file _(serialization)_ in the ONNX format only [[3]](https://www.mathworks.com/help/deeplearning/ref/exportonnxnetwork.html), so conversion is necessary to use MATLAB models with other frameworks.\n",
        "\n",
        "**TensorFlow** provides support for three different types of non-mobile serialized model formats, depending on the version of TensorFlow installed:\n",
        "1. FrozenGraph .pb files _(TensorFlow 1.X only)_ [[4]](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/tools/freeze_graph.py)\n",
        "2. SavedModel directories _(TensorFlow 1.X and 2.0)_ [[5]](https://www.tensorflow.org/guide/saved_model)\n",
        "3. HDF5 .h5 files _(TensorFlow 1.X and 2.0)_ [[6]](https://www.tensorflow.org/tutorials/keras/save_and_load)\n",
        "\n",
        "The **onnx-tf conversion tool** [[7]](https://github.com/onnx/onnx-tensorflow) was created prior to the release of TensorFlow 2.0, thus converted models are provided in the FrozenGraph .pb format only. The tool currently requires TensorFlow 1.X to be installed for conversion to work correctly [[8]](https://github.com/onnx/onnx-tensorflow/issues/521). \n",
        "\n",
        "Running the cell below ensures the correct TensorFlow version is imported and also installs the onnx-tf conversion tool."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Xj94NAARt8xR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# onnx-tf was designed for TensorFlow 1.X, so force this version.\n",
        "%tensorflow_version 1.x \n",
        "import tensorflow as tf\n",
        "\n",
        "# \"!\" allows command-line input. Use pip package manager to install the \n",
        "# conversion package\n",
        "!pip install onnx-tf"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L4YwVrCjpfiE",
        "colab_type": "text"
      },
      "source": [
        "### 2.2 Converting the Model\n",
        "\n",
        "The conversion process is quite noisy, due to three types of warnings:\n",
        "1. Deprecation warnings for functions in TensorFlow 1.x not supported in TensorFlow 2.x\n",
        "2. onnx-tf \"Fail to get since_version\" warnings.\n",
        "3. \"Unknown op\" warnings for two operators in the model graph: ConstantFill and ImageScaler.\n",
        "\n",
        "Warning **type 1** and **type 2** are harmless and can be suppressed [[9]](https://github.com/tensorflow/tensorflow/issues/27023) [[10]](https://github.com/onnx/onnx-tensorflow/issues/246). \n",
        "\n",
        "Warning **type 3** refers to experimental operators that are no longer included in operator set \"9\", but are still in use by MATLAB's ONNX converter, which uses the operator set \"8\" by default. This warning arises for two operators: **ConstantFill** [[11]](https://github.com/onnx/onnx/pull/1434) and **ImageScaler** [[12]](https://github.com/onnx/models/issues/76). These operators are still supported in conversion for backwards compatibility purposes [[13]](https://github.com/onnx/models/issues/76#issuecomment-498327977) [[14]](https://github.com/microsoft/onnxruntime/blob/master/docs/Versioning.md#Backwards-compatibility), so this warning type is not critical as of writing this notebook. In the future, however, when operator set \"8\" is no longer supported by ONNX, this notebook may become out of date. Using operator set \"9\" in MATLAB could potentially prevent warning type 3. \n",
        "\n",
        "With the warnings having been made clear, the conversion process itself requires only a few commands."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oXwuHawWvXmk",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import onnx\n",
        "from onnx_tf.backend import prepare\n",
        "\n",
        "onnx_model = onnx.load(onnx_path)\n",
        "tf_rep = prepare(onnx_model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5sMNhm4Y_pZV",
        "colab_type": "text"
      },
      "source": [
        "These lines of code produce **tf_rep**, which is a python class containing four attributes: \n",
        "1. tf_rep.graph\n",
        "2. tf_rep.inputs\n",
        "3. tf_rep.outputs\n",
        "4. tf_rep.tensor_dict\n",
        "\n",
        "These attributes can be used to identify input/output nodes, run inference, and export the intermediate model to a .pb file. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tq9KsXiKBBmn",
        "colab_type": "text"
      },
      "source": [
        "### 2.3 Exporting the Model to a .pb File\n",
        "\n",
        "Now that a tf_rep variable has been created, the converted model can be exported to a .pb file and stored within this notebook."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vOd67kNM2AOY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "pb_path = \"hasCircularShape_chartObjects_googlenet.pb\"\n",
        "tf_rep.export_graph(pb_path)\n",
        "\n",
        "assert os.path.exists(pb_path)\n",
        "print(\".pb model converted successfully.\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zoWIbZWwxrDv",
        "colab_type": "text"
      },
      "source": [
        "If you would like to save the converted model to your local disk, please run the cell below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bocn7p0fDZme",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from google.colab import files\n",
        "\n",
        "files.download(pb_path)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ij7Exq04D4An",
        "colab_type": "text"
      },
      "source": [
        "# 3 TensorFlow FrozenGraph (.pb) -> TensorFlow Lite (.tflite)\n",
        "\n",
        "Now a .pb model has been stored within the notebook, it can be prepared for Android/iOS deployment using the .tflite model format."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uZ-iKIlAGFgq",
        "colab_type": "text"
      },
      "source": [
        "### 3.1 Background Information\n",
        "\n",
        "**TensorFlow Lite** is an open source deep learning framework for on-device inference [[15]](https://www.tensorflow.org/lite). It consists of two main components: the TensorFlow Lite Converter and the TensorFlow Lite Interpreter.\n",
        "\n",
        "The **TensorFlow Lite converter** is a tool accessible using a Python API that converts trained TensorFlow models into the TensorFlow Lite format (.tflite) [[16]](https://www.tensorflow.org/lite/guide/get_started#2_convert_the_model_format). TensorFlow Lite serializes model data using the open source FlatBuffer format, which has many advantages for mobile applications [[17]](https://google.github.io/flatbuffers/).\n",
        "\n",
        "The **TensorFlow Lite interpreter** is a library that allows for inference to be run using converted TensorFlow lite models [[18]](https://www.tensorflow.org/lite/guide/get_started#3_run_inference_with_the_model). The interpreter works across multiple platforms and provides an API for running TensorFlow Lite models using Java, Swift, Objective-C, C++, and Python. Thus, a converted model can be evaluated within this notebook using the Python API, and later deployed using an API more suited for Android/iOS development.\n",
        "\n",
        "> **Note:** The FrozenGraph format (.pb) is supported by TensorFlow 1.X versions only. TensorFlow 1.X code can still be used in TensorFlow 2.0, however. If migrating the code below to a TensorFlow 2.0 environment, backwards compatibility can be ensured by modifying any deprecated function calls [[19]](https://www.tensorflow.org/guide/migrate)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wC6AVy57GP0x",
        "colab_type": "text"
      },
      "source": [
        "### 3.2 Converting the Model\n",
        "\n",
        "To use the TFLite converter to convert a FrozenGraph (.pb) file, the input and output nodes of the graph must be explicitly specified. The names of these nodes can be accessed easily using the existing tf_rep object created in **Section 2**."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kfgqYB_sUPqa",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "input_nodes = tf_rep.inputs\n",
        "output_nodes = tf_rep.outputs\n",
        "print(\"The names of the input nodes are: {}\".format(input_nodes))\n",
        "print(\"The names of the output nodes are: {}\".format(output_nodes))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sUkqNsBJUqkw",
        "colab_type": "text"
      },
      "source": [
        "With this information, the TFLiteConverter class can now be called, producing a **tflite_rep** variable which contains converted model data serialized in the TFLite FlatBuffer format."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BDNz2uWDOo1T",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "converter = tf.lite.TFLiteConverter.from_frozen_graph(pb_path,\n",
        "                                                      input_arrays=input_nodes,\n",
        "                                                      output_arrays=output_nodes)\n",
        "tflite_rep = converter.convert()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "7_1Pxu96XKpU"
      },
      "source": [
        "### 3.3 Exporting the Model to a .tflite File\n",
        "\n",
        "Now that a tflite_rep variable has been created, the converted model can be exported to a .tflite file and stored within this notebook."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "iBz2vbs-XKpd",
        "colab": {}
      },
      "source": [
        "tflite_path = \"hasCircularShape_chartObjects_googlenet.tflite\"\n",
        "open(tflite_path, \"wb\").write(tflite_rep)\n",
        "\n",
        "assert os.path.exists(tflite_path)\n",
        "print(\".tflite model converted successfully.\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "5Oe3AzzXXKp0"
      },
      "source": [
        "If you would like to save the converted model to your local disk, please run the cell below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "8cBYWkm4XKp4",
        "colab": {}
      },
      "source": [
        "from google.colab import files\n",
        "\n",
        "files.download(tflite_path)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9BJCgoNVFWcV",
        "colab_type": "text"
      },
      "source": [
        "# 4 Validation of Converted Models\n",
        "\n",
        "Now that the conversion process has finished, and various model files are stored within this notebook, the validation process can begin."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tyxRflAMX2v2",
        "colab_type": "text"
      },
      "source": [
        "### 4.1 Switching from TensorFlow 1.x to 2.x\n",
        "\n",
        "While **TensorFlow 1.x** was needed for the conversion process, **TensorFlow 2.x** will be better supported as time goes on. Thus, the remaining sections will use code that is compatible with TensorFlow 2.x versions. This ensures that further steps will be more in line with current documentation, and should make reusing this code easier.\n",
        "\n",
        "To switch versions from TensorFlow 1.x to TensorFlow 2.x within this notebook, please first restart the Colab Notebook runtime. This can be done by selecting `Runtime -> Restart runtime...` in the upper menu, as shown in the following image.\n",
        "\n",
        "> ![Restart Runtime](https://i.imgur.com/O9uZ95H.png)\n",
        "\n",
        "Once completed, please run the following cell to re-import the correct version of TensorFlow."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rJ81ixpSPSIq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%tensorflow_version 2.x\n",
        "import tensorflow as tf"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iDsGqktg-tiz",
        "colab_type": "text"
      },
      "source": [
        "> **Note:** Writing the following sections largely involved finding documentation, tutorials and StackOverflow questions written for TensorFlow 1.0, then creating a functional equivalent for that code in TensorFlow 2.0 using TensorFlow's migration guide [[19]](https://www.tensorflow.org/guide/migrate). This was challenging! If, in the future, the **onnx-tf** conversion tool allows for export in a non-\".pb\" format, then it would allow for working with TensorFlow 2.0 from the outset, which could be preferable. This need has been acknowledged in recent activity on the onnx-tensorflow GitHub page [[20]](https://github.com/onnx/onnx-tensorflow/pull/531). "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qx74mzYZFew6",
        "colab_type": "text"
      },
      "source": [
        "### 4.2 Initializing Files and Variables\n",
        "\n",
        "The .zip archive uploaded in **Section 1** of this notebook must first be extracted to a directory within the notebook workspace. This can be done using the cell below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Sk-tZjaNQNPU",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from zipfile import ZipFile\n",
        "from glob import glob\n",
        "import os\n",
        "\n",
        "# Extract the images to a folder within the workspace\n",
        "img_zip_path = 'test_images.zip'\n",
        "img_dir_path = 'img/'\n",
        "if not os.path.exists(img_dir_path):\n",
        "    with ZipFile(img_zip_path, 'r') as zip_ref:\n",
        "        zip_ref.extractall(img_dir_path)\n",
        "\n",
        "# Check that the .zip archive contains all 50 images\n",
        "TOTAL_IMAGES = 50\n",
        "assert len(glob(img_dir_path + '*/*.png')) == TOTAL_IMAGES\n",
        "print(\"The .zip was successfully extracted, and contains the required images.\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F4O_HQDXadkv",
        "colab_type": "text"
      },
      "source": [
        "Additionally, some constant values must be specified which define how the images are processed.\n",
        "\n",
        "\n",
        "* `CLASS_NAMES`: This will be used to convert a class name string (taken from the directory an image is in) into a vector of True/False values, which represents the image's label.\n",
        "* `IMG_WIDTH`, `IMG_HEIGHT`: Specifies image resizing dimensions.\n",
        "* `NUM_CHANNELS`: Specifies whether input images should be considered as grayscale or RGB."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tdb2SiunAmay",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import numpy as np\n",
        "CLASS_NAMES = np.sort(np.array([os.path.basename(path) \n",
        "                                for path in glob(img_dir_path+'/*')]))\n",
        "\n",
        "IMG_WIDTH = 224\n",
        "IMG_HEIGHT = 224\n",
        "NUM_CHANNELS = 3"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-s-8GlTjiErh",
        "colab_type": "text"
      },
      "source": [
        "### 4.3 Loading Images into a Dataset\n",
        "\n",
        "Images will be loaded into an object of the TensorFlow class **Dataset** [[21]](https://www.tensorflow.org/api_docs/python/tf/data/Dataset). The Dataset class contains methods such as **batch**, **shuffle**, and **filter** which make it easier to manipulate a larger dataset. Dataset objects are used heavily within TensorFlow tutorials and guides, such as ones for building an input pipeline for training and testing [[22]](https://www.tensorflow.org/guide/data).\n",
        "\n",
        "First, a Dataset object is created which contains the filepath of each image.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nmjaSp34_LtE",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "path_ds = tf.data.Dataset.list_files(str(img_dir_path+'*/*'))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wIW_etYyO0_v",
        "colab_type": "text"
      },
      "source": [
        "Then, a custom function is created to load an image from its filepath and preprocess it. \n",
        "\n",
        "> **Note:** Within the function below is the use of the **transpose** operation. For TensorFlow and Python packages such as matplotlib, images are defined as arrays of size `(3, height, width)`, where 3 is the number of channels. However, within MATLAB, images are defined as arrays of size `(height, width, 3)`. For loaded images to be used as input to the MATLAB model, a transpose operation is needed to shape the arrays into the correct form."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_N6fifbMO7wI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def process_image_path(image_path):\n",
        "    # Read and preprocess the image\n",
        "    img = tf.io.read_file(image_path)\n",
        "    img = tf.image.decode_png(img, channels=NUM_CHANNELS)\n",
        "    img = tf.image.resize(img, [IMG_WIDTH, IMG_HEIGHT])\n",
        "    img = tf.transpose(img, perm=[2, 0, 1])  # (224, 224, 3) -> (3, 224, 224)\n",
        "\n",
        "    # Extract the class name from the directory\n",
        "    label_name = tf.strings.split(image_path, os.path.sep)[-2]\n",
        "    # Compare the string to the list of classes to get a True/False label vector\n",
        "    label_bool = (label_name == CLASS_NAMES)\n",
        "    # Change shape from (2,) array to (2, 1) array\n",
        "    label_bool = tf.expand_dims(label_bool, axis=0)\n",
        "\n",
        "    return img, label_bool"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BBjZBnLxPCXr",
        "colab_type": "text"
      },
      "source": [
        "The Dataset method **map** is then used to apply the custom function, such that items within the filepath Dataset are mapped to an image/label Dataset. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S6vGzmJrPCgG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "labeled_ds = path_ds.map(process_image_path)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aZssNch9nBl2",
        "colab_type": "text"
      },
      "source": [
        "To demonstrate that the images have been imported correctly, the following cell displays all 50 images in a 5x10 grid."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PdIYGIu9B-lT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import matplotlib.pyplot as plt\n",
        "\n",
        "def show_dataset(ds):\n",
        "    plt.figure(figsize=(24,12))\n",
        "\n",
        "    # Iterate through subset of images\n",
        "    ds_iterator = iter(ds)\n",
        "    for n in range(TOTAL_IMAGES):\n",
        "        # Returns (1, 3, 224, 224) image tensor and (1, 2) label tensor\n",
        "        x, y = next(ds_iterator)\n",
        "\n",
        "        # Convert (1, 2) label tensor into (2) label array\n",
        "        y = np.squeeze(y.numpy())\n",
        "\n",
        "        # Convert (1, 3, 224, 224) tensor into (224, 224, 3) image\n",
        "        x = np.transpose(np.squeeze(x.numpy()), [1, 2, 0])\n",
        "\n",
        "        # Scale image to convert from [0, 255] to [0, 1]\n",
        "        x = x/255\n",
        "\n",
        "        # Plot image with its label\n",
        "        ax = plt.subplot(5,10,n+1)\n",
        "        plt.imshow(x)\n",
        "        plt.title(CLASS_NAMES[y==1][0].title())\n",
        "        plt.axis('off')\n",
        "\n",
        "show_dataset(labeled_ds)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C3RjanvwFt7I",
        "colab_type": "text"
      },
      "source": [
        "### 4.4 Validation of FrozenGraph .pb Model\n",
        "\n",
        "To run inference on a FrozenGraph, the TensorFlow 2.0 migration guide [[19]](https://www.tensorflow.org/guide/migrate#a_graphpb_or_graphpbtxt) recommends wrapping the entire graph in a `concrete_function`. The TensorFlow documentation is unclear on what a `concrete_function` is _(the link within the documentation results in a 404 Not Found error)_, but in practice doing this turns the graph into a **callable function for running inference**.\n",
        "\n",
        "To do this, three steps are needed:\n",
        "1. Loading the FrozenGraph .pb file into a `graph_def` variable."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ASOZ2MaZcfNe",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "pb_path = \"hasCircularShape_chartObjects_googlenet.pb\"\n",
        "graph_def = tf.compat.v1.GraphDef()\n",
        "loaded = graph_def.ParseFromString(open(pb_path,'rb').read())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q56PVQvftUQi",
        "colab_type": "text"
      },
      "source": [
        "2. Specifying a custom function for wrapping the loaded frozen graph. *(This is taken directly from the TensorFlow 2.0 migration guide [[19]](https://www.tensorflow.org/guide/migrate#a_graphpb_or_graphpbtxt).)*"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6VvkUDnMcmv3",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def wrap_frozen_graph(graph_def, inputs, outputs):\n",
        "    def _imports_graph_def():\n",
        "        tf.compat.v1.import_graph_def(graph_def, name=\"\")\n",
        "\n",
        "    wrapped_import = tf.compat.v1.wrap_function(_imports_graph_def, [])\n",
        "    import_graph = wrapped_import.graph\n",
        "\n",
        "    return wrapped_import.prune(\n",
        "        tf.nest.map_structure(import_graph.as_graph_element, inputs),\n",
        "        tf.nest.map_structure(import_graph.as_graph_element, outputs))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vFhW0vlnt2J0",
        "colab_type": "text"
      },
      "source": [
        "3. Finall, using this function to wrap the graph. It requires explicitly naming the input and output nodes, which were found during model conversion in **Section 2**. The \":0\" appended to the name specifies which output is desired [[23]](https://stackoverflow.com/questions/40925652/in-tensorflow-whats-the-meaning-of-0-in-a-variables-name). This becomes relevant for nodes with more than one output, but is not relevant here."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IR62QRjZuMpp",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "if loaded:\n",
        "    pb_func = wrap_frozen_graph(graph_def, inputs='data:0', outputs='softmax:0')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5pZ2Euk1uXnk",
        "colab_type": "text"
      },
      "source": [
        "The function `pb_func` can be called while iterating through the labeled dataset created earlier. It returns the values within the output node, which in this case is a **softmax** layer. The final result, `y_likelihoods`, is a list of likelihood values for each image in the dataset. This list is the same shape as `y_labels`, the list of label values determined as part of the creation of the dataset."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HYYWoMXzmSZc",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "y_likelihoods_pb, y_labels_pb = [], []\n",
        "for x, y in labeled_ds:\n",
        "    y_softmax_pb = pb_func(x)  # Run inference\n",
        "\n",
        "    y_likelihoods_pb.extend(y_softmax_pb.numpy())  # Get arrays from tensors\n",
        "    y_labels_pb.extend(y.numpy())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g-GMzdIqC6HX",
        "colab_type": "text"
      },
      "source": [
        "The predicted labels and actual labels can be determined from these lists using the **argmax** function. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eTzQOrWVpx5G",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import numpy as np\n",
        "\n",
        "y_predictions_pb = np.argmax(y_likelihoods_pb, axis=1)\n",
        "y_actual_pb = np.argmax(y_labels_pb, axis=1)\n",
        "\n",
        "print(\"Predicted labels: {}\".format(y_predictions_pb))\n",
        "print(\"Actual labels: {}\".format(y_actual_pb))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YUOG-QYyEGQj",
        "colab_type": "text"
      },
      "source": [
        "Finally, evaluation metrics can be calculated from the pair of label lists."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C-t7V4e7RVs5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "TP = tf.math.count_nonzero(y_predictions_pb * y_actual_pb)\n",
        "TN = tf.math.count_nonzero((y_predictions_pb - 1) * (y_actual_pb - 1))\n",
        "FP = tf.math.count_nonzero(y_predictions_pb * (y_actual_pb - 1))\n",
        "FN = tf.math.count_nonzero((y_predictions_pb - 1) * y_actual_pb)\n",
        "\n",
        "precision_pb = TP / (TP + FP)\n",
        "recall_pb = TP / (TP + FN)\n",
        "f1_pb = 2 * precision_pb * recall_pb / (precision_pb + recall_pb)\n",
        "\n",
        "print(\"Precision: {}, Recall: {}, F1 Score: {}\".format(precision_pb, \n",
        "                                                       recall_pb, \n",
        "                                                       f1_pb))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zJ3h-XayF68_",
        "colab_type": "text"
      },
      "source": [
        "### 4.5 Validation of TensorFlow Lite .tflite Model\n",
        "\n",
        "Running inference on the .tflite model is straightforward, and done using the **TensorFlow Lite Interpreter** first introduced in **Section 3.1**. \n",
        "\n",
        "The steps shown below _(using the Python API)_ do not vary substantially from what would be used in the C++, Java, Swift, or Objective-C APIs [[24]](https://www.tensorflow.org/lite/guide/inference#load_and_run_a_model_in_c) [[25]](https://www.tensorflow.org/lite/guide/inference#load_and_run_a_model_in_java). The only significant distinction between the Python code below and a theoretical mobile implentation would be the loading and preprocessing of input images from the mobile device's local storage.\n",
        "\n",
        "The first step for running inference using the .tflite model involves creating an interpreter instance from the model path."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "w6pQc9Z1FcLD",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "tflite_path = \"hasCircularShape_chartObjects_googlenet.tflite\"\n",
        "interpreter = tf.lite.Interpreter(tflite_path)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i3829gugKC_R",
        "colab_type": "text"
      },
      "source": [
        "Next, memory is allocated for the input and output tensors."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iLmcr4-VKDET",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "interpreter.allocate_tensors()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7hviJRK_GjHX",
        "colab_type": "text"
      },
      "source": [
        "Following this, the index of the model's input and output nodes are extracted from the interpreter. These are used to feed input images into the model, and save likelihoods from the output node."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Wmx8k1GqGiGu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "input_details = interpreter.get_input_details()\n",
        "output_details = interpreter.get_output_details()\n",
        "\n",
        "input_index = input_details[0]['index']\n",
        "output_index = output_details[0]['index']"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HLvWZ6taGgje",
        "colab_type": "text"
      },
      "source": [
        "As with running inference on the previous .pb model, the same labeled dataset is iterated through. There are some distinctions in this case, however:\n",
        "* The input image must be reshaped from a `(3, 224, 224)` array to a `(1, 3, 224, 224)` array. This may just be a quirk of the TensorFlow Lite Interpreter, as the .pb interpreter did not have this requirement.\n",
        "* The interpreter requires explicitly setting the input, invoking the interpreter, and getting the output. For the .pb inference, these steps were bundled together into a single function call.\n",
        "* The numpy array variable `x` must be cleared before the interpreter is invoked. This is done using `del(x)`. \n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ilhl807Y0fFe",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "y_likelihoods_tflite, y_labels_tflite = [], []\n",
        "\n",
        "for x, y in labeled_ds:\n",
        "    # (3, 224, 224) -> (1, 3, 224, 224)\n",
        "    x = tf.expand_dims(x, 0)\n",
        "\n",
        "    # Explicitly set input tensor\n",
        "    interpreter.set_tensor(input_index, x)\n",
        "  \n",
        "    # Free up numpy reference to internal tensor, then invoke\n",
        "    del(x)\n",
        "    interpreter.invoke()\n",
        "\n",
        "    # Explicitly get the value stored within the output tensor\n",
        "    y_softmax_tflite = interpreter.get_tensor(output_index)\n",
        "\n",
        "    y_likelihoods_tflite.extend(y_softmax_tflite)\n",
        "    y_labels_tflite.extend(y.numpy())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ezlkEHr8E8SB",
        "colab_type": "text"
      },
      "source": [
        "The remaining cells for evaluating the inference results are identical to those found within **Section 4.4**.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5PsPYv3u3eex",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import numpy as np\n",
        "\n",
        "y_predictions_tflite = np.argmax(y_likelihoods_tflite, axis=1)\n",
        "y_actual_tflite = np.argmax(y_labels_tflite, axis=1)\n",
        "\n",
        "print(\"Predicted labels: {}\".format(y_predictions_tflite))\n",
        "print(\"Actual labels: {}\".format(y_actual_tflite))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "L0j3kLKe4CDz",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "TP = tf.math.count_nonzero(y_predictions_tflite * y_actual_tflite)\n",
        "TN = tf.math.count_nonzero((y_predictions_tflite - 1) * (y_actual_tflite - 1))\n",
        "FP = tf.math.count_nonzero(y_predictions_tflite * (y_actual_tflite - 1))\n",
        "FN = tf.math.count_nonzero((y_predictions_tflite - 1) * y_actual_tflite)\n",
        "\n",
        "precision_tflite = TP / (TP + FP)\n",
        "recall_tflite = TP / (TP + FN)\n",
        "f1_tflite = 2 * precision_tflite * recall_tflite / (precision_tflite + recall_tflite)\n",
        "\n",
        "print(\"Precision: {}, Recall: {}, F1 Score: {}\".format(precision_tflite, \n",
        "                                                       recall_tflite, \n",
        "                                                       f1_tflite))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "li8oXGjUc7VP",
        "colab_type": "text"
      },
      "source": [
        "# 5 Future Work\n",
        "\n",
        "Now that the .tflite file has been validated as producing the expected inference results, the model can be further optimized, or incorporated into an Android/iOS application."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qn1HMaJGdZnb",
        "colab_type": "text"
      },
      "source": [
        "### 5.1 Further Optimization\n",
        "\n",
        "TensorFlow provides additional tools for further optimizing the model for computation requirements and disk usage [[26]](https://www.tensorflow.org/lite/guide/get_started#4_optimize_your_model_optional)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DbzMl8aNdaIk",
        "colab_type": "text"
      },
      "source": [
        "### 5.2 Integration into Mobile Application\n",
        "\n",
        "TensorFlow provides two QuickStart guides for integrating .tflite models into Android [[27]](https://www.tensorflow.org/lite/guide/android) and iOS [[28]](https://www.tensorflow.org/lite/guide/ios)."
      ]
    }
  ]
}