retrain_ssdlite_mobiledet_qat_tf1.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "license"
      },
      "source": [
        "##### *Copyright 2021 Google LLC*\n",
        "*Licensed under the Apache License, Version 2.0 (the \"License\")*"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "both",
        "id": "rKwqeqWBXANA"
      },
      "outputs": [],
      "source": [
        "# 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",
        "# https://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."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H_KfH3kOHRjL"
      },
      "source": [
        "# Retrain SSDLite Mobiledet object detector for the Coral Edge TPU (TF1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LvEMJSafnyEC"
      },
      "source": [
        "This tutorial shows you how to perform transfer-learning with a pre-trained SSDLite MobileDet model so it can detect cats and dogs. We'll use TensorFlow 1.15 for training, and then use quantization-aware training and the Edge TPU Compiler to make the model compatible with the [Coral Edge TPU](https://coral.ai/products/).\n",
        "\n",
        "Here's an example of the training results:\n",
        "\n",
        "\u003cimg src=\"https://storage.googleapis.com/site_and_emails_static_assets/Images/detection-cats-and-dogs.jpg\" /\u003e\n",
        "\n",
        "**Beware:** This Colab takes about 3 hours to finish and that's only because we dialed-down a couple training parameters so it doesn't take even longer (which would cause your Colab instance to time-out and stop your training) or run out of memory. Although the results are still good, you'll get even better results (a more accurate model) if you [connect this notebook to a local runtime](https://research.google.com/colaboratory/local-runtimes.html), or upgrade to [Colab Pro](https://colab.research.google.com/signup) (which provides faster GPUs, longer run-times, and more memory), and then increase the number of training steps and the batch size."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "viewin-badges"
      },
      "source": [
        "\u003ca href=\"https://colab.research.google.com/github/google-coral/tutorials/blob/master/retrain_ssdlite_mobiledet_qat_tf1.ipynb\" target=\"_parent\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open in Colab\"\u003e\u003c/a\u003e\n",
        "\u0026nbsp;\u0026nbsp;\u0026nbsp;\u0026nbsp;\n",
        "\u003ca href=\"https://github.com/google-coral/tutorials/blob/master/retrain_ssdlite_mobiledet_qat_tf1.ipynb\" target=\"_parent\"\u003e\u003cimg src=\"https://img.shields.io/static/v1?logo=GitHub\u0026label=\u0026color=333333\u0026style=flat\u0026message=View%20on%20GitHub\" alt=\"View in GitHub\"\u003e\u003c/a\u003e\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8i9X8Ssj3rlw"
      },
      "source": [
        "To start running all the code in this tutorial, select **Runtime \u003e Run all** in the Colab toolbar."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IIIFhejc3yE3"
      },
      "source": [
        "## Install the code dependencies"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZBZXusuSSVES"
      },
      "outputs": [],
      "source": [
        "! pip install -U numpy==1.19.5\n",
        "! pip install -U pycocotools==2.0.1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2tPji39Y9dOP"
      },
      "source": [
        "Make sure you're using TensorFlow 1.15:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qNW-c3MD9Vks"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "  # This %tensorflow_version magic only works in Colab.\n",
        "  %tensorflow_version 1.x\n",
        "except Exception:\n",
        "  pass\n",
        "# For your non-Colab code, be sure you have tensorflow==1.15\n",
        "import tensorflow as tf\n",
        "assert tf.__version__.startswith('1')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kyuTvP3p9hIa"
      },
      "source": [
        "Then we need to build the [TF1 Object Detection API](https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf1.md):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H1d6HUYkMBTg"
      },
      "outputs": [],
      "source": [
        "! pip install tf_slim\n",
        "! git clone https://github.com/tensorflow/models.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SnSBXH75XMcD"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "os.environ['PYTHONPATH'] += ':/content/models/research/'\n",
        "os.environ['PYTHONPATH'] += ':/content/models/research/slim/'\n",
        "os.environ['PYTHONPATH'] += ':/content/models/research/object_detection/utils/'\n",
        "os.environ['PYTHONPATH'] += ':/content/models/research/object_detection'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AzUbTD7f5T6A"
      },
      "outputs": [],
      "source": [
        "! apt-get install protobuf-compiler"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Yhaz9nH55c_r"
      },
      "outputs": [],
      "source": [
        "%cd models/research\n",
        "# Compile all the protobuf dependencies\n",
        "! protoc object_detection/protos/*.proto --python_out=.\n",
        "# Set up and install the object detection API\n",
        "! cp object_detection/packages/tf1/setup.py .\n",
        "! python -m pip install .\n",
        "# Run a test to make sure setup is correct\n",
        "! python object_detection/builders/model_builder_test.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ECIrLo5O9HUy"
      },
      "source": [
        "## Prepare the training dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b5DSZPig3laE"
      },
      "source": [
        "Now let's get the training dataset. We're using is the [Oxford-IIIT Pets Dataset](https://www.robots.ox.ac.uk/~vgg/data/pets/), which includes 37 different classes of dog and cat breeds. But for this tutorial, we'll use just 2 classes for training: Abyssinian cats and American Bulldogs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wI-7dJozsupa"
      },
      "outputs": [],
      "source": [
        "# Download the dataset files\n",
        "%mkdir /content/dataset\n",
        "%cd /content/dataset\n",
        "! wget http://www.robots.ox.ac.uk/~vgg/data/pets/data/images.tar.gz\n",
        "! wget http://www.robots.ox.ac.uk/~vgg/data/pets/data/annotations.tar.gz\n",
        "! tar zxf images.tar.gz\n",
        "! tar zxf annotations.tar.gz"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZSpk4nRTHAIY"
      },
      "outputs": [],
      "source": [
        "# Only picking Abyssinian Cat and the American Bulldog\n",
        "# If you wish to train the model on all classes, simply skip this entire cell\n",
        "! cp /content/dataset/annotations/list.txt /content/dataset/annotations/list_petsdataset.txt\n",
        "! cp /content/dataset/annotations/trainval.txt /content/dataset/annotations/trainval_petsdataset.txt\n",
        "! cp /content/dataset/annotations/test.txt /content/dataset/annotations/test_petsdataset.txt\n",
        "! grep \"Abyssinian\" /content/dataset/annotations/list_petsdataset.txt \u003e  /content/dataset/annotations/list.txt\n",
        "! grep \"american_bulldog\" /content/dataset/annotations/list_petsdataset.txt \u003e\u003e /content/dataset/annotations/list.txt\n",
        "! grep \"Abyssinian\" /content/dataset/annotations/trainval_petsdataset.txt \u003e /content/dataset/annotations/trainval.txt\n",
        "! grep \"american_bulldog\" /content/dataset/annotations/trainval_petsdataset.txt \u003e\u003e /content/dataset/annotations/trainval.txt\n",
        "! grep \"Abyssinian\" /content/dataset/annotations/test_petsdataset.txt \u003e /content/dataset/annotations/test.txt\n",
        "! grep \"american_bulldog\" /content/dataset/annotations/test_petsdataset.txt \u003e\u003e /content/dataset/annotations/test.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zXL3dxtx8McO"
      },
      "source": [
        "Now we convert the images to the [TFRecord file format](https://www.tensorflow.org/tutorials/load_data/tfrecord) we need for training:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qdF1oqYWuL4R"
      },
      "outputs": [],
      "source": [
        "%cd /content/models/research\n",
        "! cp object_detection/data/pet_label_map.pbtxt /content/dataset\n",
        "! python3 object_detection/dataset_tools/create_pet_tf_record.py \\\n",
        "    --label_map_path=\"/content/dataset/pet_label_map.pbtxt\" \\\n",
        "    --data_dir=\"/content/dataset\" \\\n",
        "    --output_dir=\"/content/dataset\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6CxGNKU0715P"
      },
      "source": [
        "Also, create a labels file for our classes:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Lnacf_1A_Hco"
      },
      "outputs": [],
      "source": [
        "OUTPUT_DIR = '/content/output_ssdlite_mobiledet_dog_vs_cat'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sEXzXbth71l6"
      },
      "outputs": [],
      "source": [
        "! mkdir $OUTPUT_DIR\n",
        "! echo \"0 Abyssinian\" \u003e\u003e \"$OUTPUT_DIR/labels.txt\"\n",
        "! echo \"1 american_bulldog\" \u003e\u003e \"$OUTPUT_DIR/labels.txt\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8cCnG9Sh6DKJ"
      },
      "source": [
        "## Prepare the model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mZllwseL-y1h"
      },
      "source": [
        "First, let's download the pre-trained MobileDet model that's optimized for the Edge TPU:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f5v347uRtKzF"
      },
      "outputs": [],
      "source": [
        "!mkdir /content/pretrained_model\n",
        "%cd /content/pretrained_model\n",
        "! wget http://download.tensorflow.org/models/object_detection/ssdlite_mobiledet_edgetpu_320x320_coco_2020_05_19.tar.gz\n",
        "! tar xvf ssdlite_mobiledet_edgetpu_320x320_coco_2020_05_19.tar.gz"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MGR65PTY-17o"
      },
      "source": [
        "Then we define a few variables for our training parameters:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UAt1xs4BH1dM"
      },
      "outputs": [],
      "source": [
        "# Ideally, we'd use more steps, but much larger will reach the system timeout for a free Colab environment\n",
        "# If you have Colab Pro or you're running offline, try 25000\n",
        "NUM_STEPS = 10000\n",
        "\n",
        "# Ideally, batch size would be larger, but smaller is necessary to avoid OOM Killer on free Colab environments\n",
        "# If you have Colab Pro or you're running offline, try 64\n",
        "BATCH_SIZE = 32\n",
        "\n",
        "# For this tutorial, we're training just two classes\n",
        "# If you train with the whole cats/dogs dataset, it's 37 classes\n",
        "NUM_CLASSES = 2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a_Pvcans__YW"
      },
      "source": [
        "Next, we grab the [object detection config file](https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/configuring_jobs.md) corresponding to the pre-trained model we're using, and then make some adjustments to match our training pipeline, such as the training dataset, the quantization settings, and other training parameters.\n",
        "\n",
        "**Tip:** When training your own detection model, you should also consider adjusting the [anchor box parameters](https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/configuring_jobs.md#anchor-box-parameters), which can improve the model accuracy and training speed. To generate new anchor box parameters for this dog/cat dataset, check out the Colab to [Generate SSD anchor box aspect ratios using k-means\n",
        "clustering](https://colab.sandbox.google.com/github/tensorflow/models/blob/master/research/object_detection/colab_tutorials/generate_ssd_anchor_box_aspect_ratios_using_k_means_clustering.ipynb). Then use the `.config` file output by that notebook in the following section (update `config_path`).\n",
        " "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Mg1C8UwStK7i"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "from google.protobuf import text_format\n",
        "from object_detection.protos import pipeline_pb2\n",
        "import os\n",
        "\n",
        "pipeline = pipeline_pb2.TrainEvalPipelineConfig()                                                                                                                                                                                                          \n",
        "config_path = '/content/models/research/object_detection/samples/configs/ssdlite_mobiledet_edgetpu_320x320_coco_sync_4x4.config'\n",
        "with tf.gfile.GFile(config_path, \"r\") as f:                                                                                                                                                                                                                     \n",
        "    proto_str = f.read()                                                                                                                                                                                                                                          \n",
        "    text_format.Merge(proto_str, pipeline)\n",
        "\n",
        "pipeline.train_input_reader.tf_record_input_reader.input_path[:] = ['/content/dataset/pet_faces_train.record-?????-of-00010']\n",
        "pipeline.train_input_reader.label_map_path = '/content/dataset/pet_label_map.pbtxt'\n",
        "pipeline.eval_input_reader[0].tf_record_input_reader.input_path[:] = ['/content/dataset/pet_faces_val.record-?????-of-00010']\n",
        "pipeline.eval_input_reader[0].label_map_path = '/content/dataset/pet_label_map.pbtxt'\n",
        "pipeline.train_config.fine_tune_checkpoint = '/content/pretrained_model/ssdlite_mobiledet_edgetpu_320x320_coco_2020_05_19/fp32/model.ckpt'\n",
        "pipeline.train_config.batch_size = BATCH_SIZE\n",
        "pipeline.train_config.num_steps = NUM_STEPS\n",
        "pipeline.model.ssd.num_classes = NUM_CLASSES\n",
        "# Enable ssdlite, this should already be enabled in the config we downloaded, but this is just to make sure.\n",
        "pipeline.model.ssd.box_predictor.convolutional_box_predictor.kernel_size = 3\n",
        "pipeline.model.ssd.box_predictor.convolutional_box_predictor.use_depthwise = True\n",
        "pipeline.model.ssd.feature_extractor.use_depthwise = True\n",
        "# Quantization Aware Training\n",
        "pipeline.graph_rewriter.quantization.delay = 0\n",
        "pipeline.graph_rewriter.quantization.weight_bits = 8\n",
        "pipeline.graph_rewriter.quantization.activation_bits = 8\n",
        "\n",
        "config_text = text_format.MessageToString(pipeline)                                                                                                                                                                                                        \n",
        "with tf.gfile.Open(config_path, \"wb\") as f:                                                                                                                                                                                                                       \n",
        "    f.write(config_text)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QUNpzaVWGRG7"
      },
      "source": [
        "You can now see our pipeline changes in the config file (such as `num_classes: 2`):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e_3xrVUbGPqV"
      },
      "outputs": [],
      "source": [
        "! cat /content/models/research/object_detection/samples/configs/ssdlite_mobiledet_edgetpu_320x320_coco_sync_4x4.config"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J9nvA1YUHot1"
      },
      "source": [
        "## Launch TensorBoard"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vg0TJs8oHwMo"
      },
      "source": [
        "TensorBoard is optional but provides very helpful visualizations of your training progress and accuracy evaluations.\n",
        "\n",
        "Because TensorBoard runs as a webserver on your local machine—and we're actually running this on a Colab virtual environment—we'll use a tool called [ngrok](https://ngrok.com/) to make this server accessible with a public URL:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "whIM6gZz2DEn"
      },
      "outputs": [],
      "source": [
        "%cd /content\n",
        "! wget https://bin.equinox.io/c/4VmDzA7iaHb/ngrok-stable-linux-amd64.zip\n",
        "! unzip -o ngrok-stable-linux-amd64.zip"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XvJ9X0qBfSIX"
      },
      "outputs": [],
      "source": [
        "# Starts tensorboard, so we can monitor the training process.\n",
        "get_ipython().system_raw(\n",
        "    'tensorboard --logdir {} --host 0.0.0.0 --port 6006 \u0026'\n",
        "    .format('/content/train')\n",
        ")\n",
        "get_ipython().system_raw('./ngrok http 6006 \u0026')\n",
        "print('Click this link to view training progress in TensorBoard:')\n",
        "import time\n",
        "time.sleep(1)\n",
        "!curl -s http://localhost:4040/api/tunnels | python3 -c \\\n",
        "    \"import sys, json; print(json.load(sys.stdin)['tunnels'][0]['public_url'])\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V9nAbMhtN1yM"
      },
      "source": [
        "## Train the model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "czY8o-mSKUQw"
      },
      "source": [
        "Here's the part that takes about 3 hours to finish (we included a training stopwatch)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oaBn6FZANYXg"
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\n",
        "start = datetime.now()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Idw8OgZQtLKG"
      },
      "outputs": [],
      "source": [
        "%cd /content/models/research/\n",
        "! python3 object_detection/model_main.py \\\n",
        "    --logtostderr=true \\\n",
        "    --model_dir=/content/train \\\n",
        "    --pipeline_config_path=/content/models/research/object_detection/samples/configs/ssdlite_mobiledet_edgetpu_320x320_coco_sync_4x4.config"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KgHc5octNtYT"
      },
      "outputs": [],
      "source": [
        "end = datetime.now()\n",
        "duration = end - start\n",
        "seconds_in_hour = 60 * 60\n",
        "hours, seconds = divmod(duration.seconds, seconds_in_hour)\n",
        "minutes = int(seconds / 60)\n",
        "print('TRAINING TIME:', str(hours) + ':' + str(minutes if minutes \u003e 10 else '%02d' % minutes))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2mjvbKt3N6JO"
      },
      "source": [
        "## Export the full model\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZSBKL7y-aZ1a"
      },
      "source": [
        "This gives us the inference\n",
        "graph, associated checkpoint files, a frozen inference graph, and a\n",
        "SavedModel:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KicA9DZZmnKe"
      },
      "outputs": [],
      "source": [
        "! python3 /content/models/research/object_detection/export_inference_graph.py \\\n",
        "    --input_type=image_tensor \\\n",
        "    --pipeline_config_path=/content/models/research/object_detection/samples/configs/ssdlite_mobiledet_edgetpu_320x320_coco_sync_4x4.config \\\n",
        "    --output_directory=/content/inference_graph \\\n",
        "    --trained_checkpoint_prefix=/content/train/model.ckpt-$NUM_STEPS"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7N-DMHQaavsS"
      },
      "source": [
        "### Evaluate the full model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gRRXF0TrMqHW"
      },
      "outputs": [],
      "source": [
        "# Let's download some test data from flickr.\n",
        "! mkdir /content/test\n",
        "! cd /content/test\n",
        "! wget https://live.staticflickr.com/7921/46683787864_86c9501c24_c_d.jpg -O /content/test/image1.jpg\n",
        "! wget https://live.staticflickr.com/4/8451898_8bedb2ae53_c_d.jpg -O /content/test/image2.jpg\n",
        "! wget https://live.staticflickr.com/2654/3997966238_f454845087_c_d.jpg -O /content/test/image3.jpg\n",
        "! wget https://live.staticflickr.com/2818/34032378096_5309537c9f_c_d.jpg -O /content/test/image4.jpg\n",
        "! wget https://live.staticflickr.com/8682/28214087384_4c7711584d_c_d.jpg -O /content/test/image5.jpg"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mhg-2OA-m-9n"
      },
      "outputs": [],
      "source": [
        "# Do a Quick Evaluation on the inference graph model.\n",
        "import numpy as np\n",
        "import os\n",
        "import sys\n",
        "import tensorflow as tf\n",
        "\n",
        "from collections import defaultdict\n",
        "from matplotlib import pyplot as plt\n",
        "from PIL import Image\n",
        "\n",
        "from object_detection.utils import ops as utils_ops\n",
        "from object_detection.utils import label_map_util\n",
        "from object_detection.utils import visualization_utils as vis_util\n",
        "%matplotlib inline\n",
        "\n",
        "# Initialize tf.Graph()\n",
        "detection_graph = tf.Graph()\n",
        "with detection_graph.as_default():\n",
        "  od_graph_def = tf.GraphDef()\n",
        "  with tf.gfile.GFile('/content/inference_graph/frozen_inference_graph.pb', 'rb') as fid:\n",
        "    serialized_graph = fid.read()\n",
        "    od_graph_def.ParseFromString(serialized_graph)\n",
        "    tf.import_graph_def(od_graph_def, name='')\n",
        "\n",
        "# Loads labels\n",
        "label_map = label_map_util.load_labelmap('/content/dataset/pet_label_map.pbtxt')\n",
        "categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=2, use_display_name=True)\n",
        "category_index = label_map_util.create_category_index(categories)\n",
        "\n",
        "# Run Inference and populates results in a dict.\n",
        "def run_inference(graph, image):\n",
        "  with graph.as_default():\n",
        "    with tf.Session() as sess:\n",
        "      ops = tf.get_default_graph().get_operations()\n",
        "      all_tensor_names = [output.name for op in ops for output in op.outputs]\n",
        "      tensor_dict = {}\n",
        "      tensor_keys = ['num_detections', 'detection_boxes', 'detection_scores', 'detection_classes']\n",
        "      for key in tensor_keys:\n",
        "        tensor_name = key + ':0'\n",
        "        if tensor_name in all_tensor_names:\n",
        "          tensor_dict[key] = tf.get_default_graph().get_tensor_by_name(tensor_name)\n",
        "      \n",
        "      # Actual inference.\n",
        "      image_tensor = tf.get_default_graph().get_tensor_by_name('image_tensor:0')\n",
        "      output_dict = sess.run(tensor_dict, feed_dict={image_tensor: np.expand_dims(image, 0)})\n",
        "\n",
        "      output_dict['num_detections'] = int(output_dict['num_detections'][0])\n",
        "      output_dict['detection_classes'] = output_dict['detection_classes'][0].astype(np.uint8)\n",
        "      output_dict['detection_boxes'] = output_dict['detection_boxes'][0]\n",
        "      output_dict['detection_scores'] = output_dict['detection_scores'][0]\n",
        "  return output_dict\n",
        "\n",
        "test_image_path = [os.path.join('/content/test', 'image{}.jpg'.format(i)) for i in range(1, 6)]\n",
        "for image_path in test_image_path:\n",
        "  print('Evaluating:', image_path)\n",
        "  image = Image.open(image_path)\n",
        "  img_width, img_height = image.size\n",
        "  image_np = np.array(image.getdata()).reshape((img_height, img_width, 3)).astype(np.uint8)\n",
        "  # Run inference.\n",
        "  output_dict = run_inference(detection_graph, image_np)\n",
        "  # Visualization of the results of a detection.\n",
        "  vis_util.visualize_boxes_and_labels_on_image_array(\n",
        "      image_np,\n",
        "      output_dict['detection_boxes'],\n",
        "      output_dict['detection_classes'],\n",
        "      output_dict['detection_scores'],\n",
        "      category_index,\n",
        "      use_normalized_coordinates=True,\n",
        "      line_thickness=8)\n",
        "  plt.figure(figsize=(12, 8))\n",
        "  plt.imshow(image_np)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r6etUm4ya5Lc"
      },
      "source": [
        "## Export to TF Lite"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6sHP33iZBlil"
      },
      "outputs": [],
      "source": [
        "# Export this model to tflite_graph format\n",
        "%cd /content/models/research\n",
        "\n",
        "! python3 object_detection/export_tflite_ssd_graph.py \\\n",
        "  --pipeline_config_path=/content/models/research/object_detection/samples/configs/ssdlite_mobiledet_edgetpu_320x320_coco_sync_4x4.config \\\n",
        "  --trained_checkpoint_prefix=/content/train/model.ckpt-$NUM_STEPS \\\n",
        "  --output_directory=$OUTPUT_DIR \\\n",
        "  --add_postprocessing_op=true"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-NYcwbhdEmBN"
      },
      "outputs": [],
      "source": [
        "# Convert to a tflite file (for CPU)\n",
        "! tflite_convert \\\n",
        "  --output_file=\"$OUTPUT_DIR/ssdlite_mobiledet_dog_vs_cat.tflite\" \\\n",
        "  --graph_def_file=\"$OUTPUT_DIR/tflite_graph.pb\" \\\n",
        "  --inference_type=QUANTIZED_UINT8 \\\n",
        "  --input_arrays=\"normalized_input_image_tensor\" \\\n",
        "  --output_arrays=\"TFLite_Detection_PostProcess,TFLite_Detection_PostProcess:1,TFLite_Detection_PostProcess:2,TFLite_Detection_PostProcess:3\" \\\n",
        "  --mean_values=128 \\\n",
        "  --std_dev_values=128 \\\n",
        "  --input_shapes=1,320,320,3 \\\n",
        "  --allow_custom_ops"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tFKaHfQFcItu"
      },
      "source": [
        "### Evaluate the TF Lite model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_A-IYRpm-Ate"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import re\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "from PIL import ImageDraw\n",
        "%matplotlib inline\n",
        "\n",
        "# Creates tflite interpreter\n",
        "interpreter = tf.lite.Interpreter(OUTPUT_DIR + '/ssdlite_mobiledet_dog_vs_cat.tflite')\n",
        "interpreter.allocate_tensors()\n",
        "interpreter.invoke() # warmup\n",
        "input_details = interpreter.get_input_details()\n",
        "output_details = interpreter.get_output_details()\n",
        "width = input_details[0]['shape'][2]\n",
        "height = input_details[0]['shape'][1]\n",
        "\n",
        "def read_label_file(file_path):\n",
        "  with open(file_path, 'r', encoding='utf-8') as f:\n",
        "    lines = f.readlines()\n",
        "  ret = {}\n",
        "  for row_number, content in enumerate(lines):\n",
        "    pair = re.split(r'[:\\s]+', content.strip(), maxsplit=1)\n",
        "    if len(pair) == 2 and pair[0].strip().isdigit():\n",
        "      ret[int(pair[0])] = pair[1].strip()\n",
        "    else:\n",
        "      ret[row_number] = content.strip()\n",
        "  return ret\n",
        "\n",
        "def run_inference(interpreter, image):\n",
        "  interpreter.set_tensor(input_details[0]['index'], image)\n",
        "  interpreter.invoke()\n",
        "  boxes = interpreter.get_tensor(output_details[0]['index'])[0]\n",
        "  classes = interpreter.get_tensor(output_details[1]['index'])[0]\n",
        "  scores = interpreter.get_tensor(output_details[2]['index'])[0]\n",
        "  # num_detections = interpreter.get_tensor(output_details[3]['index'])[0]\n",
        "  return boxes, classes, scores\n",
        "\n",
        "test_image_paths = [os.path.join('/content/test', 'image{}.jpg'.format(i)) for i in range(1, 6)]\n",
        "for image_path in test_image_paths:\n",
        "  print('Evaluating:', image_path)\n",
        "  image = Image.open(image_path)\n",
        "  image_width, image_height = image.size\n",
        "  draw = ImageDraw.Draw(image)\n",
        "  resized_image = image.resize((width, height))\n",
        "  np_image = np.asarray(resized_image)\n",
        "  input_tensor = np.expand_dims(np_image, axis=0)\n",
        "  # Run inference\n",
        "  boxes, classes, scores = run_inference(interpreter, input_tensor)\n",
        "  # Draw results on image\n",
        "  colors = {0:(128, 255, 102), 1:(102, 255, 255)}\n",
        "  labels = read_label_file(OUTPUT_DIR + '/labels.txt')\n",
        "  for i in range(len(boxes)):\n",
        "    if scores[i] \u003e .7:\n",
        "      ymin = int(max(1, (boxes[i][0] * image_height)))\n",
        "      xmin = int(max(1, (boxes[i][1] * image_width)))\n",
        "      ymax = int(min(image_height, (boxes[i][2] * image_height)))\n",
        "      xmax = int(min(image_width, (boxes[i][3] * image_width)))\n",
        "      draw.rectangle((xmin, ymin, xmax, ymax), width=7, outline=colors[int(classes[i])])\n",
        "      draw.rectangle((xmin, ymin, xmax, ymin-10), fill=colors[int(classes[i])])\n",
        "      text = labels[int(classes[i])] + ' ' + str(scores[i]*100) + '%'\n",
        "      draw.text((xmin+2, ymin-10), text, fill=(0,0,0), width=2)\n",
        "  display(image)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X7HrTCSMcYay"
      },
      "source": [
        "## Compile it for the Edge TPU"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7t7TRJ5gcnS1"
      },
      "source": [
        "Install the [Edge TPU Compiler](https://coral.ai/docs/edgetpu/compiler/):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jUy2l-pJFUfV"
      },
      "outputs": [],
      "source": [
        "! curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add -\n",
        "! echo \"deb https://packages.cloud.google.com/apt coral-edgetpu-stable main\" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list\n",
        "! sudo apt-get update\n",
        "! sudo apt-get install edgetpu-compiler"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vj-omWPscj6P"
      },
      "source": [
        "Compile the model for the Edge TPU:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OW45-68aFbNd"
      },
      "outputs": [],
      "source": [
        "%cd $OUTPUT_DIR\n",
        "\n",
        "! edgetpu_compiler -s ssdlite_mobiledet_dog_vs_cat.tflite"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uVuB1XImcvsv"
      },
      "source": [
        "## Package and download all files"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Xs4zdU1wzpsB"
      },
      "outputs": [],
      "source": [
        "! cp -r /content/train/model.ckpt-$NUM_STEPS* $OUTPUT_DIR\n",
        "! cp -r /content/inference_graph/* $OUTPUT_DIR\n",
        "\n",
        "%cd /content/\n",
        "! tar cvf output_ssdlite_mobiledet_dog_vs_cat.tar.gz output_ssdlite_mobiledet_dog_vs_cat"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b13V2T5wfG5I"
      },
      "outputs": [],
      "source": [
        "from google.colab import files\n",
        "files.download('/content/output_ssdlite_mobiledet_dog_vs_cat.tar.gz')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_TZTwG7nhm0C"
      },
      "source": [
        "## Run it on the Edge TPU\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RwywT4ZpQjLf"
      },
      "source": [
        "You can now run the model on your [Coral device](https://www.coral.ai/products) with acceleration on the Edge TPU.\n",
        "\n",
        "First, find some new photos to try. Remember that you've trained this model to recognize just two classes: Abyssinian cats and\n",
        "American Bulldogs. So here are a couple images that should provide results (provided by the\n",
        "[Open Images Dataset](https://storage.googleapis.com/openimages/web/index.html)): \n",
        "\n",
        "```\n",
        "wget https://c4.staticflickr.com/8/7580/15865399370_ffa5b49d20_z.jpg -O dog.jpg \u0026\u0026 \\\n",
        "wget https://c6.staticflickr.com/9/8534/8652503705_687d957a29_z.jpg -O cat.jpg\n",
        "```\n",
        "\n",
        "Then, try running an inference using [this example code for the PyCoral API](https://github.com/google-coral/pycoral/blob/master/examples/detect_image.py). Just clone that repo and run the script using the model files you downloaded above (also be sure you have [installed the PyCoral API](https://coral.ai/software/#pycoral-api)):\n",
        "\n",
        "```\n",
        "git clone https://github.com/google-coral/pycoral\n",
        "\n",
        "cd pycoral/examples/\n",
        "\n",
        "python3 detect_image.py \\\n",
        "  --model ssdlite_mobiledet_dog_vs_cat_edgetpu.tflite \\\n",
        "  --labels labels.txt \\\n",
        "  --input dog.jpg \\\n",
        "  --output dog_result.jpg\n",
        "```\n",
        "\n",
        "Check out more examples for running inference at [coral.ai/examples](https://coral.ai/examples/#code-examples/)."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "license"
      ],
      "machine_shape": "hm",
      "name": "Retrain SSDLite MobileDet detector for the Edge TPU (TF1)",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}