retrain_efficientdet_model_maker_tf2.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": "Gb7qyhNL1yWt"
      },
      "source": [
        "# Retrain EfficientDet for the Edge TPU with TensorFlow Lite Model Maker"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sr3q-gvm3cI8"
      },
      "source": [
        "In this tutorial, we'll retrain the EfficientDet-Lite object detection model (derived from [EfficientDet](https://ai.googleblog.com/2020/04/efficientdet-towards-scalable-and.html)) using the [TensorFlow Lite Model Maker library](https://www.tensorflow.org/lite/guide/model_maker), and then compile it to run on the [Coral Edge TPU](https://www.coral.ai/products/). All in about 30 minutes.\n",
        "\n",
        "By default, we'll retrain the model using a publicly available dataset of salad photos, teaching the model to recognize a salad and some of the ingredients. But we've also provided code so you can upload your own training dataset in the Pascal VOC XML format.\n",
        "\n",
        "Here's an example of the salad training results:\n",
        "\n",
        "\u003cimg src=\"https://storage.googleapis.com/site_and_emails_static_assets/Images/efficientdet-salads.png?\" width=\"400\" hspace=\"0\"\u003e\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zHZyMPbMsAns"
      },
      "source": [
        "\u003ca href=\"https://colab.research.google.com/github/google-coral/tutorials/blob/master/retrain_efficientdet_model_maker_tf2.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_efficientdet_model_maker_tf2.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": "PzoNZRp4sVxK"
      },
      "source": [
        "If you want to run the notebook with the salad dataset, you can run the whole thing now by clicking **Runtime \u003e Run all** in the Colab toolbar. But if you want to use your own dataset, then continue down to [Load the training data](#scrollTo=H0XM-oIfhgQ7) and follow the instructions there.\n",
        "\n",
        "**Note:** If using a custom dataset, beware that if your dataset includes more than 20 classes, you'll probably have slower inference speeds compared to if you have fewer classes. This is due to an aspect of the EfficientDet architecture in which a certain layer cannot compile for the Edge TPU when it carries more than 20 classes."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2vvAObmTqglq"
      },
      "source": [
        "## Import the required packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qhl8lqVamEty"
      },
      "outputs": [],
      "source": [
        "!pip install -q tflite-model-maker"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XtxiUeZEiXpt"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import os\n",
        "\n",
        "from tflite_model_maker.config import ExportFormat\n",
        "from tflite_model_maker import model_spec\n",
        "from tflite_model_maker import object_detector\n",
        "\n",
        "import tensorflow as tf\n",
        "assert tf.__version__.startswith('2')\n",
        "\n",
        "tf.get_logger().setLevel('ERROR')\n",
        "from absl import logging\n",
        "logging.set_verbosity(logging.ERROR)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H0XM-oIfhgQ7"
      },
      "source": [
        "## Load the training data\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QusnAYjwM3uq"
      },
      "source": [
        "To use the default salad training dataset, just run all the code below as-is.\n",
        "\n",
        "But if you want to train with your own image dataset, follow these steps:\n",
        "\n",
        "1. Be sure your dataset is annotated in Pascal VOC XML (various tools can help create VOC annotations, such as [LabelImg](https://github.com/tzutalin/labelImg#labelimg)). Then create a ZIP file with all your JPG images and XML files (JPG and XML files can all be in one directory or in separate directories).\n",
        "2. Click the **Files** tab in the left panel and just drag-drop your ZIP file there to upload it.\n",
        "3. Use the following drop-down option to set **`use_custom_dataset`** to True.\n",
        "4. If your dataset is already split into separate directories for training, validation, and testing, also set **`dataset_is_split`** to True. (If your dataset is not split, leave it False and we'll split it below.)\n",
        "5. Then skip to [Load your own Pascal VOC dataset](#scrollTo=ZljJ25RAnj5x) and follow the rest of the instructions there.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "wVQ2VQu7MuQa"
      },
      "outputs": [],
      "source": [
        "use_custom_dataset = False #@param [\"False\", \"True\"] {type:\"raw\"}\n",
        "\n",
        "dataset_is_split = False #@param [\"False\", \"True\"] {type:\"raw\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-9VoZ_Rud2V3"
      },
      "source": [
        "### Load the salads CSV dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BRd13bfetO7B"
      },
      "source": [
        "\n",
        "\n",
        "Model Maker requires that we load our dataset using the [`DataLoader`](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/object_detector/DataLoader) API. So in this case, we'll load it from a CSV file that defines 175 images for training, 25 images for validation, and 25 images for testing.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "04ObtdneqvP5"
      },
      "outputs": [],
      "source": [
        "if not use_custom_dataset:\n",
        "  train_data, validation_data, test_data = object_detector.DataLoader.from_csv('gs://cloud-ml-data/img/openimage/csv/salads_ml_use.csv')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7W0Njyxoq0gH"
      },
      "source": [
        "If you want to load your own dataset as a CSV file, you can learn more about the format in [Formatting a training data CSV](https://cloud.google.com/vision/automl/object-detection/docs/csv-format). You can load your CSV either from [Cloud Storage](https://cloud.google.com/storage) (as shown above) or from a local path.\n",
        "\n",
        "[`DataLoader`](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/object_detector/DataLoader) can also load your dataset in other formats, such as from a set of TFRecord files or from a local directory using the Pascal VOC format (shown below for a custom dataset). "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZljJ25RAnj5x"
      },
      "source": [
        "### (Optional) Load your own Pascal VOC dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ei5BahmPn_wR"
      },
      "source": [
        "To use your custom dataset, you need to modify a few variables here, such as your ZIP filename, your label map, and the path to your images/annotations:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mz_suhWiqc7A"
      },
      "outputs": [],
      "source": [
        "if use_custom_dataset:\n",
        "\n",
        "  # The ZIP file you uploaded:\n",
        "  !unzip dataset.zip\n",
        "\n",
        "  # Your labels map as a dictionary (zero is reserved):\n",
        "  label_map = {1: 'apple', 2: 'banana'} \n",
        "\n",
        "  if dataset_is_split:\n",
        "    # If your dataset is already split, specify each path:\n",
        "    train_images_dir = 'dataset/train/images'\n",
        "    train_annotations_dir = 'dataset/train/annotations'\n",
        "    val_images_dir = 'dataset/validation/images'\n",
        "    val_annotations_dir = 'dataset/validation/annotations'\n",
        "    test_images_dir = 'dataset/test/images'\n",
        "    test_annotations_dir = 'dataset/test/annotations'\n",
        "  else:\n",
        "    # If it's NOT split yet, specify the path to all images and annotations\n",
        "    images_in = 'dataset/images'\n",
        "    annotations_in = 'dataset/annotations'\n",
        "  "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d8Rh3IWRb0xw"
      },
      "source": [
        "Now you're ready to train the model with your custom dataset. But before you run the notebook, you should also skip to the [Export to TensorFlow Lite](#scrollTo=_yB_XMpqGlLs) section and change the `TFLITE_FILENAME` and `LABLES_FILENAME` for your exported files.\n",
        "\n",
        "Then run the whole notebook by clicking **Runtime \u003e Run all**."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "C35meprE8xzf"
      },
      "outputs": [],
      "source": [
        "#@markdown Be sure you run this cell. It's hiding the `split_dataset()` function used in the next code block.\n",
        "\n",
        "import os\n",
        "import random\n",
        "import shutil\n",
        "\n",
        "def split_dataset(images_path, annotations_path, val_split, test_split, out_path):\n",
        "  \"\"\"Splits a directory of sorted images/annotations into training, validation, and test sets.\n",
        "\n",
        "  Args:\n",
        "    images_path: Path to the directory with your images (JPGs).\n",
        "    annotations_path: Path to a directory with your VOC XML annotation files,\n",
        "      with filenames corresponding to image filenames. This may be the same path\n",
        "      used for images_path.\n",
        "    val_split: Fraction of data to reserve for validation (float between 0 and 1).\n",
        "    test_split: Fraction of data to reserve for test (float between 0 and 1).\n",
        "  Returns:\n",
        "    The paths for the split images/annotations (train_dir, val_dir, test_dir)\n",
        "  \"\"\"\n",
        "  _, dirs, _ = next(os.walk(images_path))\n",
        "\n",
        "  train_dir = os.path.join(out_path, 'train')\n",
        "  val_dir = os.path.join(out_path, 'validation')\n",
        "  test_dir = os.path.join(out_path, 'test')\n",
        "\n",
        "  IMAGES_TRAIN_DIR = os.path.join(train_dir, 'images')\n",
        "  IMAGES_VAL_DIR = os.path.join(val_dir, 'images')\n",
        "  IMAGES_TEST_DIR = os.path.join(test_dir, 'images')\n",
        "  os.makedirs(IMAGES_TRAIN_DIR, exist_ok=True)\n",
        "  os.makedirs(IMAGES_VAL_DIR, exist_ok=True)\n",
        "  os.makedirs(IMAGES_TEST_DIR, exist_ok=True)\n",
        "\n",
        "  ANNOT_TRAIN_DIR = os.path.join(train_dir, 'annotations')\n",
        "  ANNOT_VAL_DIR = os.path.join(val_dir, 'annotations')\n",
        "  ANNOT_TEST_DIR = os.path.join(test_dir, 'annotations')\n",
        "  os.makedirs(ANNOT_TRAIN_DIR, exist_ok=True)\n",
        "  os.makedirs(ANNOT_VAL_DIR, exist_ok=True)\n",
        "  os.makedirs(ANNOT_TEST_DIR, exist_ok=True)\n",
        "\n",
        "  # Get all filenames for this dir, filtered by filetype\n",
        "  filenames = os.listdir(os.path.join(images_path))\n",
        "  filenames = [os.path.join(images_path, f) for f in filenames if (f.endswith('.jpg'))]\n",
        "  # Shuffle the files, deterministically\n",
        "  filenames.sort()\n",
        "  random.seed(42)\n",
        "  random.shuffle(filenames)\n",
        "  # Get exact number of images for validation and test; the rest is for training\n",
        "  val_count = int(len(filenames) * val_split)\n",
        "  test_count = int(len(filenames) * test_split)\n",
        "  for i, file in enumerate(filenames):\n",
        "    source_dir, filename = os.path.split(file)\n",
        "    annot_file = os.path.join(annotations_path, filename.replace(\"jpg\", \"xml\"))\n",
        "    if i \u003c val_count:\n",
        "      shutil.copy(file, IMAGES_VAL_DIR)\n",
        "      shutil.copy(annot_file, ANNOT_VAL_DIR)\n",
        "    elif i \u003c val_count + test_count:\n",
        "      shutil.copy(file, IMAGES_TEST_DIR)\n",
        "      shutil.copy(annot_file, ANNOT_TEST_DIR)\n",
        "    else:\n",
        "      shutil.copy(file, IMAGES_TRAIN_DIR)\n",
        "      shutil.copy(annot_file, ANNOT_TRAIN_DIR)\n",
        "  return (train_dir, val_dir, test_dir)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KWROlVNA54xZ"
      },
      "outputs": [],
      "source": [
        "# We need to instantiate a separate DataLoader for each split dataset\n",
        "if use_custom_dataset:\n",
        "  if dataset_is_split:\n",
        "    train_data = object_detector.DataLoader.from_pascal_voc(\n",
        "        train_images_dir, train_annotations_dir, label_map=label_map)\n",
        "    validation_data = object_detector.DataLoader.from_pascal_voc(\n",
        "        val_images_dir, val_annotations_dir, label_map=label_map)\n",
        "    test_data = object_detector.DataLoader.from_pascal_voc(\n",
        "        test_images_dir, test_annotations_dir, label_map=label_map)\n",
        "  else:\n",
        "    train_dir, val_dir, test_dir = split_dataset(images_in, annotations_in,\n",
        "                                                 val_split=0.2, test_split=0.2,\n",
        "                                                 out_path='split-dataset')\n",
        "    train_data = object_detector.DataLoader.from_pascal_voc(\n",
        "        os.path.join(train_dir, 'images'),\n",
        "        os.path.join(train_dir, 'annotations'), label_map=label_map)\n",
        "    validation_data = object_detector.DataLoader.from_pascal_voc(\n",
        "        os.path.join(val_dir, 'images'),\n",
        "        os.path.join(val_dir, 'annotations'), label_map=label_map)\n",
        "    test_data = object_detector.DataLoader.from_pascal_voc(\n",
        "        os.path.join(test_dir, 'images'),\n",
        "        os.path.join(test_dir, 'annotations'), label_map=label_map)\n",
        "    \n",
        "  print(f'train count: {len(train_data)}')\n",
        "  print(f'validation count: {len(validation_data)}')\n",
        "  print(f'test count: {len(test_data)}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S8clx0KPutCM"
      },
      "source": [
        "## Select the model spec"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vn61LJ9QbOPi"
      },
      "source": [
        "Model Maker supports the EfficientDet-Lite family of object detection models that are compatible with the Edge TPU. (EfficientDet-Lite is derived from [EfficientDet](https://ai.googleblog.com/2020/04/efficientdet-towards-scalable-and.html), which offers state-of-the-art accuracy in a small model size). There are several model sizes you can choose from:\n",
        "\n",
        "|| Model architecture | Size(MB)* | Latency(ms)** | Average Precision*** |\n",
        "|-|--------------------|-----------|---------------|----------------------|\n",
        "|| EfficientDet-Lite0 | 5.7       | 37.4            | 30.4%               |\n",
        "|| EfficientDet-Lite1 | 7.6       | 56.3            | 34.3%               |\n",
        "|| EfficientDet-Lite2 | 10.2      | 104.6           | 36.0%               |\n",
        "|| EfficientDet-Lite3 | 14.4      | 107.6           | 39.4%               |\n",
        "| \u003ctd colspan=4\u003e\u003cbr\u003e\u003ci\u003e* File size of the compiled Edge TPU models. \u003cbr/\u003e** Latency measured on a desktop CPU with a Coral USB Accelerator. \u003cbr/\u003e*** Average Precision is the mAP (mean Average Precision) on the COCO 2017 validation dataset.\u003c/i\u003e\u003c/td\u003e |\n",
        "\n",
        "Beware that the Lite2 and Lite3 models do not fit onto the Edge TPU's onboard memory, so you'll see even greater latency when using those, due to the cost of fetching data from the host system memory. Maybe this extra latency is okay for your application, but if it's not and you require the precision of the larger models, then you can [pipeline the model across multiple Edge TPUs](https://coral.ai/docs/edgetpu/pipeline/) (more about this when we compile the model below).\n",
        "\n",
        "For this tutorial, we'll use Lite0:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SM9gePHw9Jv1"
      },
      "outputs": [],
      "source": [
        "spec = object_detector.EfficientDetLite0Spec()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rnCzdzs0-Rbo"
      },
      "source": [
        "The [`EfficientDetLite0Spec`](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/object_detector/EfficientDetLite0Spec) constructor also supports several arguments that specify training options, such as the max number of detections (default is 25 for the TF Lite model) and whether to use Cloud TPUs for training. You can also use the constructor to specify the number of training epochs and the batch size, but you can also specify those in the next step."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5qjq2UEHCLUi"
      },
      "source": [
        "## Create and train the model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2uZkLR6N6gDR"
      },
      "source": [
        "Now we need to create our model according to the model spec, load our dataset into the model, specify training parameters, and begin training. \n",
        "\n",
        "Using Model Maker, we accomplished all of that with [`create()`](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/object_detector/create):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kwlYdTcg63xy"
      },
      "outputs": [],
      "source": [
        "model = object_detector.create(train_data=train_data, \n",
        "                               model_spec=spec, \n",
        "                               validation_data=validation_data, \n",
        "                               epochs=50, \n",
        "                               batch_size=10, \n",
        "                               train_whole_model=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3n5-o3vvGfnJ"
      },
      "source": [
        "## Evaluate the model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-BzCHLWJ6h7q"
      },
      "source": [
        "Now we'll use the test dataset to evaluate how well the model performs with data it has never seen before.\n",
        "\n",
        "The [`evaluate()`](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/object_detector/ObjectDetector#evaluate) method provides output in the style of [COCO evaluation metrics](https://cocodataset.org/#detection-eval):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8xmnl6Yy7ARn"
      },
      "outputs": [],
      "source": [
        "model.evaluate(test_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PEon9xd2BDS_"
      },
      "source": [
        "Because the default batch size for [EfficientDetLite models](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/object_detector/EfficientDetSpec) is 64, this needs only 1 step to go through all 25 images in the salad test set. You can also specify the `batch_size` argument when you call [`evaluate()`](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/object_detector/ObjectDetector#evaluate)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_yB_XMpqGlLs"
      },
      "source": [
        "## Export to TensorFlow Lite"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CgCDMe0e6jlT"
      },
      "source": [
        "Next, we'll export the model to the TensorFlow Lite format. By default, the [`export()`](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/object_detector/ObjectDetector#export) method performs [full integer post-training quantization](https://www.tensorflow.org/lite/performance/post_training_quantization#full_integer_quantization), which is exactly what we need for compatibility with the Edge TPU. (Model Maker uses the same dataset we gave to our model spec as a representative dataset, which is required for full-int quantization.)\n",
        "\n",
        "We just need to specify the export directory and format. By default, it exports to TF Lite, but we also want a labels file, so we declare both:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2Cu9cxX5Qu-e"
      },
      "outputs": [],
      "source": [
        "TFLITE_FILENAME = 'efficientdet-lite-salad.tflite'\n",
        "LABELS_FILENAME = 'salad-labels.txt'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rKd6qk7TbxYO"
      },
      "outputs": [],
      "source": [
        "model.export(export_dir='.', tflite_filename=TFLITE_FILENAME, label_filename=LABELS_FILENAME,\n",
        "             export_format=[ExportFormat.TFLITE, ExportFormat.LABEL])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b94hZ-exOCRB"
      },
      "source": [
        "### Evaluate the TF Lite model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZQpahAIBqBPp"
      },
      "source": [
        "Exporting the model to TensorFlow Lite can affect the model accuracy, due to the reduced numerical precision from quantization and because the original TensorFlow model uses per-class [non-max supression (NMS)](https://www.coursera.org/lecture/convolutional-neural-networks/non-max-suppression-dvrjH) for post-processing, while the TF Lite model uses global NMS, which is faster but less accurate.\n",
        "\n",
        "Therefore you should always evaluate the exported TF Lite model and be sure it still meets your requirements:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RS3Ell_lqH4e"
      },
      "outputs": [],
      "source": [
        "model.evaluate_tflite(TFLITE_FILENAME, test_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ph88z7PdOeX7"
      },
      "source": [
        "### Try the TFLite model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "me6_RwPZqNhX"
      },
      "source": [
        "Just to be sure of things, let's run the model ourselves with an image from the test set. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eag7jTOASGFW"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "\n",
        "# If you're using a custom dataset, we take a random image from the test set:\n",
        "if use_custom_dataset:\n",
        "  images_path = test_images_dir if dataset_is_split else os.path.join(test_dir, \"images\")\n",
        "  filenames = os.listdir(os.path.join(images_path))\n",
        "  random_index = random.randint(0,len(filenames)-1)\n",
        "  INPUT_IMAGE = os.path.join(images_path, filenames[random_index])\n",
        "else:\n",
        "  # Download a test salad image\n",
        "  INPUT_IMAGE = 'salad-test.jpg'\n",
        "  DOWNLOAD_URL = \"https://storage.googleapis.com/cloud-ml-data/img/openimage/3/2520/3916261642_0a504acd60_o.jpg\"\n",
        "  !wget -q -O $INPUT_IMAGE $DOWNLOAD_URL"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZBecI78ZaxsO"
      },
      "source": [
        "To simplify our code, we'll use the [PyCoral API](https://coral.ai/docs/reference/py/):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TmgtGBqua1N3"
      },
      "outputs": [],
      "source": [
        "! python3 -m pip install --extra-index-url https://google-coral.github.io/py-repo/ pycoral"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GkXtipXKqXp4"
      },
      "outputs": [],
      "source": [
        "from PIL import Image\n",
        "from PIL import ImageDraw\n",
        "from PIL import ImageFont\n",
        "\n",
        "import tflite_runtime.interpreter as tflite \n",
        "from pycoral.adapters import common\n",
        "from pycoral.adapters import detect\n",
        "from pycoral.utils.dataset import read_label_file\n",
        "\n",
        "def draw_objects(draw, objs, scale_factor, labels):\n",
        "  \"\"\"Draws the bounding box and label for each object.\"\"\"\n",
        "  COLORS = np.random.randint(0, 255, size=(len(labels), 3), dtype=np.uint8)\n",
        "  for obj in objs:\n",
        "    bbox = obj.bbox\n",
        "    color = tuple(int(c) for c in COLORS[obj.id])\n",
        "    draw.rectangle([(bbox.xmin * scale_factor, bbox.ymin * scale_factor),\n",
        "                    (bbox.xmax * scale_factor, bbox.ymax * scale_factor)],\n",
        "                   outline=color, width=3)\n",
        "    font = ImageFont.truetype(\"LiberationSans-Regular.ttf\", size=15)\n",
        "    draw.text((bbox.xmin * scale_factor + 4, bbox.ymin * scale_factor + 4),\n",
        "              '%s\\n%.2f' % (labels.get(obj.id, obj.id), obj.score),\n",
        "              fill=color, font=font)\n",
        "\n",
        "# Load the TF Lite model\n",
        "labels = read_label_file(LABELS_FILENAME)\n",
        "interpreter = tflite.Interpreter(TFLITE_FILENAME)\n",
        "interpreter.allocate_tensors()\n",
        "\n",
        "# Resize the image for input\n",
        "image = Image.open(INPUT_IMAGE)\n",
        "_, scale = common.set_resized_input(\n",
        "    interpreter, image.size, lambda size: image.resize(size, Image.ANTIALIAS))\n",
        "\n",
        "# Run inference\n",
        "interpreter.invoke()\n",
        "objs = detect.get_objects(interpreter, score_threshold=0.4, image_scale=scale)\n",
        "\n",
        "# Resize again to a reasonable size for display\n",
        "display_width = 500\n",
        "scale_factor = display_width / image.width\n",
        "height_ratio = image.height / image.width\n",
        "image = image.resize((display_width, int(display_width * height_ratio)))\n",
        "draw_objects(ImageDraw.Draw(image), objs, scale_factor, labels)\n",
        "image"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oxgWQyYOqZha"
      },
      "source": [
        "## Compile for the Edge TPU\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A0QLiwCj9Pw6"
      },
      "source": [
        "First we need to download the Edge TPU Compiler:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Oy3QIn_YqaRP"
      },
      "outputs": [],
      "source": [
        "! curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -\n",
        "\n",
        "! echo \"deb https://packages.cloud.google.com/apt coral-edgetpu-stable main\" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list\n",
        "\n",
        "! sudo apt-get update\n",
        "\n",
        "! sudo apt-get install edgetpu-compiler"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qRWewhqFqeL_"
      },
      "source": [
        "Before compiling the `.tflite` file for the Edge TPU, it's important to consider whether your model will fit into the Edge TPU memory. \n",
        "\n",
        "The Edge TPU has approximately 8 MB of SRAM for [caching model paramaters](https://coral.ai/docs/edgetpu/compiler/#parameter-data-caching), so any model close to or over 8 MB will not fit onto the Edge TPU memory. That means the inference times are longer, because some model parameters must be fetched from the host system memory.\n",
        "\n",
        "One way to elimiate the extra latency is to use [model pipelining](https://coral.ai/docs/edgetpu/pipeline/), which splits the model into segments that can run on separate Edge TPUs in series. This can significantly reduce the latency for big models.\n",
        "\n",
        "The following table provides recommendations for the number of Edge TPUs to use with each EfficientDet-Lite model.\n",
        "\n",
        "| Model architecture | Minimum TPUs | Recommended TPUs\n",
        "|--------------------|-------|-------|\n",
        "| EfficientDet-Lite0 | 1     | 1     |\n",
        "| EfficientDet-Lite1 | 1     | 1     |\n",
        "| EfficientDet-Lite2 | 1     | 2     |\n",
        "| EfficientDet-Lite3 | 2     | 2     |\n",
        "| EfficientDet-Lite4 | 2     | 3     |\n",
        "\n",
        "If you need extra Edge TPUs for your model, then update `NUMBER_OF_TPUS` here:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LZdonJGCqieU"
      },
      "outputs": [],
      "source": [
        "NUMBER_OF_TPUS =  1\n",
        "\n",
        "!edgetpu_compiler $TFLITE_FILENAME -d --num_segments=$NUMBER_OF_TPUS"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O2CjkduY02DF"
      },
      "source": [
        "**Beware when using multiple segments:** The Edge TPU Comiler divides the model such that all segments have roughly equal amounts of parameter data, but that does not mean all segments have the same latency. Especially when dividing an SSD model such as EfficientDet, this results in a latency-imbalance between segments, because SSD models have a large post-processing op that actually executes on the CPU, not on the Edge TPU. So although segmenting your model this way is better than running the whole model on just one Edge TPU, we recommend that you segment the EfficientDet-Lite model using our [profiling-based partitioner tool](https://github.com/google-coral/libcoral/tree/master/coral/tools/partitioner#profiling-based-partitioner-for-the-edge-tpu-compiler), which measures each segment's latency on the Edge TPU and then iteratively adjusts the segmentation sizes to provide balanced latency between all segments."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KyBBvyqx0XRn"
      },
      "source": [
        "## Download the files"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M43URVgg0ZcB"
      },
      "outputs": [],
      "source": [
        "from google.colab import files\n",
        "\n",
        "files.download(TFLITE_FILENAME)\n",
        "files.download(TFLITE_FILENAME.replace('.tflite', '_edgetpu.tflite'))\n",
        "files.download(LABELS_FILENAME)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8fYwl4Rt8myF"
      },
      "source": [
        "## Run the model on the Edge TPU"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dv2Ingvx8pcI"
      },
      "source": [
        "You can now run the model with acceleration on the Edge TPU.\n",
        "\n",
        "First, download an image of a salad on your device with an Edge TPU. For example, you can use the same one we tested above: \n",
        "\n",
        "```\n",
        "wget https://storage.googleapis.com/cloud-ml-data/img/openimage/3/2520/3916261642_0a504acd60_o.jpg -O salad.jpg\n",
        "```\n",
        "\n",
        "And then make sure you have [installed the PyCoral API](https://coral.ai/software/#pycoral-api).\n",
        "\n",
        "Now run an inference using [this example code for the PyCoral API](https://github.com/google-coral/pycoral/blob/master/examples/detect_image.py). Just clone the GitHub repo and run the example, passing it the model files from this tutorial:\n",
        "\n",
        "```\n",
        "git clone https://github.com/google-coral/pycoral\n",
        "\n",
        "cd pycoral/examples/\n",
        "\n",
        "python3 detect_image.py \\\n",
        "  --model efficientdet-lite-salads_edgetpu.tflite \\\n",
        "  --labels salad-labels.txt \\\n",
        "  --input salad.jpg \\\n",
        "  --output salad_result.jpg\n",
        "```\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HS4u77W5gnzQ"
      },
      "source": [
        "## More resources\n",
        "\n",
        "* For more information about the Model Maker library used in this tutorial, see the [TensorFlow Lite Model Maker guide](https://www.tensorflow.org/lite/guide/model_maker) and [API reference](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker).\n",
        "\n",
        "* For other transfer learning tutorials that are compatible with the Edge TPU, see the [Colab tutorials for Coral](https://github.com/google-coral/tutorials#colab-tutorials-for-coral).\n",
        "\n",
        "* You can also find more examples that show how to run inference on the Edge TPU at [coral.ai/examples](https://coral.ai/examples/#code-examples/)."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "license"
      ],
      "name": "Retrain EfficientDet-Lite detector for the Edge TPU (TF2)",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}