Armadillo_IoT_G4_create_detection_model.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Armadillo-IoT_G4_create_detection_model.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "private_outputs": true,
      "authorship_tag": "ABX9TyMzeHlPqbbXtpbGiNGcJb40",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/atmark-techno/google-colab-notebooks/blob/master/Armadillo_IoT_G4_create_detection_model.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7TkdmBMqZ2xn"
      },
      "source": [
        "# はじめに\n",
        "\n",
        "本ドキュメントではArmadillo-IoT ゲートウェイ G4上で使用できるTFLite形式の推論モデルを、既存のモデルをベースに学習する転移学習を行い作成するサンプルを紹介します。\n",
        "\n",
        "「Armadillo Base OS開発ガイド」の「5.1.1 教師データの用意」まで完了していることを前提としています。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mye5iniEZ5zw"
      },
      "source": [
        "# 準備"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dl3IyEdkaHxX"
      },
      "source": [
        "## Google Driveのマウント\n",
        "\n",
        "最終的に作成するTFLite形式のファイルと、それを生成する元となるSavedModel形式のファイルはGoogle Driveに保存しますので、最初にGoogle Driveをマウントします。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hdJunZmAz3co"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('./Drive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ce9edLlEaNv9"
      },
      "source": [
        "## TensorFlowのインストール"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_AVvELPB0igi"
      },
      "source": [
        "!pip install -U --pre tensorflow==\"2.8.0\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oxS_9r0uaSDh"
      },
      "source": [
        "## modelsリポジトリのクローン\n",
        "\n",
        "TensorFlow公式が提供しているmodelsリポジトリをクローンします。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E0ziLUlW0vKE"
      },
      "source": [
        "import pathlib\n",
        "\n",
        "if \"models\" in pathlib.Path.cwd().parts:\n",
        "  while \"models\" in pathlib.Path.cwd().parts:\n",
        "    os.chdir('..')\n",
        "elif not pathlib.Path('models').exists():\n",
        "  !git clone --depth 1 https://github.com/tensorflow/models"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EW9XD8CBakIY"
      },
      "source": [
        "## Object Detection APIのインストール\n",
        "\n",
        "modelsリポジトリに含まれるObject Detection APIをインストールします。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Wd7cc5lB0xki"
      },
      "source": [
        "%%bash\n",
        "cd models/research\n",
        "protoc object_detection/protos/*.proto --python_out=.\n",
        "cp object_detection/packages/tf2/setup.py .\n",
        "python -m pip install --use-feature=2020-resolver ."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aaHf_MuJa4Ox"
      },
      "source": [
        "Object Detection APIが正しくインストールされたかテストします。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uga0TL-s0zkh"
      },
      "source": [
        "!python models/research/object_detection/builders/model_builder_tf2_test.py"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NdgCrsbzbBWO"
      },
      "source": [
        "## ライブラリ群のインポート\n",
        "\n",
        "以後必要なライブラリをまとめてインポートします。\n",
        "\n",
        "**デフォルトのままではcv2のインポートに失敗するので、先にopencv-python-headlessをインストールする必要があります。**  \n",
        "https://github.com/opencv/opencv-python/issues/591"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install \"opencv-python-headless<4.3\""
      ],
      "metadata": {
        "id": "Zo8pbPVBm1S1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YwWjL_pBUs15"
      },
      "source": [
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import os\n",
        "import random\n",
        "import io\n",
        "import imageio\n",
        "import glob\n",
        "import scipy.misc\n",
        "import shutil\n",
        "import copy\n",
        "import cv2\n",
        "import numpy as np\n",
        "from six import BytesIO\n",
        "from PIL import Image, ImageDraw, ImageFont\n",
        "from IPython.display import display, Javascript\n",
        "from IPython.display import Image as IPyImage\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "from object_detection.utils import label_map_util\n",
        "from object_detection.utils import config_util\n",
        "from object_detection.utils import visualization_utils as viz_utils\n",
        "from object_detection.utils import colab_utils\n",
        "from object_detection.utils import config_util\n",
        "from object_detection.builders import model_builder\n",
        "\n",
        "%matplotlib inline"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qGefSV3RdzYz"
      },
      "source": [
        "## ディレクトリの生成\n",
        "\n",
        "作業用のディレクトリを生成します。\n",
        "作業用ディレクトリは、Google Colaboratoryのランタイムが終了すると自動的に消去されます。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ubpS4M65d6BI"
      },
      "source": [
        "%%bash\n",
        "# 作業用ディレクトリ\n",
        "mkdir -p working/01_train_data\n",
        "mkdir -p working/02_annotations\n",
        "mkdir -p working/03_pretrained_model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S5g_xGLAfX0P"
      },
      "source": [
        "## 教師データの用意\n",
        "\n",
        "Armadillo Base OS開発ガイドの「5.1. 推論モデルの作成」の手順に従って作成した教師データを、それぞれ本Colab内の所定のディレクトリにアップロードしてください。\n",
        "\n",
        "アップロードは画面左のファイルツリーからアップロードできます。\n",
        "\n",
        "![アップロード方法](https://download.atmark-techno.com/armadillo-guide-std/document/images/howto_upload.png)\n",
        "\n",
        "1. 画像ファイル(ATDEの~/datset_gauge/downloads/gauge/*.jpg)  \n",
        "  /content/working/01_train_data/にアップロード  \n",
        "  \n",
        "  ![画像ファイルのアップロード](https://download.atmark-techno.com/armadillo-guide-std/document/images/upload_images.png)\n",
        "2. .xmlファイル及び、tf_label_map.pbtxt(ATDEの~/dataset_gauge/annotations)  \n",
        "  /content/working/02_annotations/にアップロード  \n",
        "  \n",
        "  ![画像ファイルのアップロード](https://download.atmark-techno.com/armadillo-guide-std/document/images/upload_annotations.png)\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%bash\n",
        "cp /content/Drive/MyDrive/datasets/object-detection/read_meter/downloads/gauge/*.jpg /content/working/01_train_data/\n",
        "cp /content/Drive/MyDrive/datasets/object-detection/read_meter/annotations/*.xml /content/working/02_annotations/\n",
        "cp /content/Drive/MyDrive/datasets/object-detection/read_meter/xmls/label_map.pbtxt /content/working/02_annotations/"
      ],
      "metadata": {
        "id": "NRp09ODSFo-o"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9yvoQq1-bMno"
      },
      "source": [
        "# 転移学習\n",
        "\n",
        "ここからは既存の推論モデルをベースに、検出したい物体が検出できるようにチューニングを行う、転移学習の具体的な手順の一例を紹介します。\n",
        "\n",
        "本手順では物体検出の転移学習についてのみ取り扱います。\n",
        "物体検出以外の推論モデルにつきましては、大きく手順が異なることがありますので、適宜公式のマニュアル等を参照してください。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xgFomG-jb10v"
      },
      "source": [
        "## ベースとなる推論モデルをダウンロード\n",
        "\n",
        "本ドキュメントでは、[COCOデータセット](https://cocodataset.org/#explore)でトレーニング済みのSSD MobileNet V2をベースモデルとします。\n",
        "\n",
        "ベースモデルを[TensorFlow2 Detection Model Zoo](https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf2_detection_zoo.md)からダウンロード・展開します。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zsp15tZ-ARcu"
      },
      "source": [
        "%%bash\n",
        "cd working/03_pretrained_model/\n",
        "wget http://download.tensorflow.org/models/object_detection/tf2/20200711/ssd_mobilenet_v2_320x320_coco17_tpu-8.tar.gz\n",
        "tar xf ssd_mobilenet_v2_320x320_coco17_tpu-8.tar.gz\n",
        "rm ssd_mobilenet_v2_320x320_coco17_tpu-8.tar.gz"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gdlP8PqScgZz"
      },
      "source": [
        "## 教師データの分割\n",
        "\n",
        "用意した教師データを学習用データ75%、検証用データ25%の割合で分けます。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0ZyL5t5cmIzI"
      },
      "source": [
        "# 各ディレクトリ名を定義\n",
        "images_dir = '/content/working/01_train_data'\n",
        "annotations_dir = '/content/working/02_annotations'\n",
        "\n",
        "train_dir = '/content/working/train'\n",
        "train_images_dir = '/content/working/train/gauge/JPEGImages'\n",
        "train_annotations_dir = '/content/working/train/VOC2007/Annotations'\n",
        "\n",
        "val_dir = '/content/working/val'\n",
        "val_images_dir = '/content/working/val/gauge/JPEGImages'\n",
        "val_annotations_dir = '/content/working/val/VOC2007/Annotations'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qUmfGDYPmHvO"
      },
      "source": [
        "# 定義したディレクトリを作成\n",
        "!mkdir -p $train_images_dir\n",
        "!mkdir -p $train_annotations_dir\n",
        "!mkdir -p $train_dir/VOC2007/ImageSets/Main\n",
        "\n",
        "!mkdir -p $val_images_dir\n",
        "!mkdir -p $val_annotations_dir\n",
        "!mkdir -p $val_dir/VOC2007/ImageSets/Main"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j_Q07s4ulhbe"
      },
      "source": [
        "# ファイル数をカウント\n",
        "file_count = len(glob.glob(annotations_dir + '/*.xml'))\n",
        "print('File count : ' + str(file_count))\n",
        "\n",
        "# 学習データと検証データに分割\n",
        "train_ratio = 0.75\n",
        "\n",
        "file_list = glob.glob(annotations_dir + '/*.xml')\n",
        "random_sample_list = random.sample(file_list, file_count)\n",
        "\n",
        "# ディレクトリへコピー\n",
        "for index, filepath in enumerate(random_sample_list):\n",
        "  imagepath = os.path.join(images_dir, os.path.splitext(os.path.basename(filepath))[0] + '.jpg')\n",
        "  if index < int(file_count * train_ratio):\n",
        "    # 学習データとしてコピー\n",
        "    shutil.copy2(filepath, train_annotations_dir)\n",
        "    shutil.copy2(imagepath, train_images_dir)\n",
        "  else:\n",
        "    # 検証データとしてコピー\n",
        "    shutil.copy2(filepath, val_annotations_dir)\n",
        "    shutil.copy2(imagepath, val_images_dir)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cXG6XuUG7wqP"
      },
      "source": [
        "## tfrecord形式に変換\n",
        "\n",
        "分けた教師データをそれぞれtfrecord形式に変換します。\n",
        "\n",
        "tfrecordとは、TensorFlowが推奨するTensorFlowでデータを学習させる際のフォーマットです。\n",
        "詳細は[公式ドキュメント](https://www.tensorflow.org/tutorials/load_data/tfrecord?hl=ja)をご確認ください。\n",
        "\n",
        "TensorFlowはもちろんcsvやxmlなどの他のフォーマットからでも学習可能ですが、tfrecord形式にすることで、メモリに収まりきらないような大規模なデータセットを扱うことができます。\n",
        "\n",
        "Object Detection APIには、PascalVOC形式のxmlアノテーションファイルをtfrecordに変換するスクリプトが付属しています。\n",
        "今回はこちらをそのまま利用するために、スクリプトに合わせたディレクトリ構成を構築します。\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZKnpDZ8xKFoW"
      },
      "source": [
        "# 学習用データセットをtfrecord化\n",
        "%%bash\n",
        "python models/research/object_detection/dataset_tools/create_pascal_tf_record.py \\\n",
        "--data_dir=/content/working/train \\\n",
        "--year=VOC2007 \\\n",
        "--output_path=/content/working/train.record \\\n",
        "--label_map_path=/content/working/02_annotations/label_map.pbtxt"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kz-I3BaeOKmc"
      },
      "source": [
        "# 検証用データセットをtfrecord化\n",
        "%%bash\n",
        "python models/research/object_detection/dataset_tools/create_pascal_tf_record.py \\\n",
        "--data_dir=/content/working/val \\\n",
        "--set=val \\\n",
        "--year=VOC2007 \\\n",
        "--output_path=/content/working/eval.record \\\n",
        "--label_map_path=/content/working/02_annotations/label_map.pbtxt"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WLVB-PTOjhzE"
      },
      "source": [
        "## pipeline.configの修正\n",
        "\n",
        "/content/working/03_pretrained_model/ssd_mobilenet_v2_320x320_coco17_tpu-8/pipeline.configを修正します。\n",
        "\n",
        "上記のpipeline.configへのリンクをクリックすると、Google Colaboratory上で編集することが可能です。\n",
        "\n",
        "以下の箇所を修正してください。\n",
        "\n",
        "* 3行目: num_classes: 90 → 1\n",
        " * 検出する物の種類数です。\n",
        " * 今回は\"gauge\"のみなので1とします。\n",
        "* 138行目: batch_size: 512 → 16\n",
        " * 確率的勾配降下法を使用する場合、異常値の影響を小さくするためにデータセットをいくつかのサブセットに分けて学習します。\n",
        " * そのサブセットに含まれるデータの数を定義します。\n",
        " * 大きくすると学習時の負荷は高まり、環境によってはメモリ不足でエラーとなることがあります。\n",
        " * 2のn乗の数値を設定してください。\n",
        "* 162行目: fine_tune_checkpoint: \"PATH_TO_BE_CONFIGURED\" → \"/content/working/03_pretrained_model/ssd_mobilenet_v2_320x320_coco17_tpu-8/checkpoint/ckpt-0\"\n",
        " * 転移学習のベースモデルを指定します。\n",
        " * \"ckpt-0\"までを記載すれば大丈夫です。\n",
        "* 168行目: fine_tune_checkpoint_type: \"classification\" → \"detection\"\n",
        " * このパラメータによって、学習済みモデルのどの層を学習し直すかを選択できます。\n",
        " * このパラメータは、以下の3つから選べます。\n",
        "   * \"classification\"  \n",
        "   学習済みモデルから、分類バックボーン部分以外を学習し直します。主に、学習済み画像分類モデルから物体検出モデルを学習する際によく使用されます。\n",
        "   * \"detection\"  \n",
        "   学習済みモデルから、ボックス予測ヘッドとクラス予測ヘッドを学習し直します。主に、学習済み物体検出モデルに新たにクラスを追加する際によく使用されます。\n",
        "   * \"full\"  \n",
        "   学習済みモデルから、全てのパラメータを学習し直します。今回は使用しませんが、取り扱いが複雑なパラメータなので使用する際には注意してください。\n",
        "  * 今回は学習済み物体検出モデルに新たに\"gauge\"クラスを追加したいので、\"detection\"に設定します。\n",
        "* 172行目: label_map_path: \"PATH_TO_BE_CONFIGURED\" → \"/content/working/02_annotations/label_map.pbtxt\"\n",
        " * 学習用にラベルマップファイル(label_map.pbtxt)のパスを指定します。\n",
        "* 174行目: input_path: \"PATH_TO_BE_CONFIGURED\" → \"/content/working/train.record\"\n",
        " * 学習用データ(train.record)のパスを指定します。\n",
        "* 182行目: label_map_path: \"PATH_TO_BE_CONFIGURED\" → \"/content/working/02_annotations/label_map.pbtxt\"\n",
        " * 検証用にラベルマップファイル(label_map.pbtxt)のパスを指定します。\n",
        "* 186行目: input_path: \"PATH_TO_BE_CONFIGURED\" → \"/content/working/eval.record\"\n",
        " * 検証用データ(train.record)のパスを指定します。\n",
        "\n",
        "編集後は、Ctrl+Sで保存できます。\n",
        "\n",
        "pipeline.configについては[公式ドキュメント](https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/configuring_jobs.md)も合わせて参照してください。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TQU3icVZXse3"
      },
      "source": [
        "## 学習結果保存用ディレクトリ作成\n",
        "\n",
        "Google Drive内にArmadillo-BaseOS-Guideという名前で、転移学習後の推論モデルを保存するディレクトリを作成します。\n",
        "\n",
        "既に同名のディレクトリが存在すると消去して作り直すので注意してください。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BNkCdJbuYKsd"
      },
      "source": [
        "!rm -rf Drive/MyDrive/Armadillo-BaseOS-Guide/\n",
        "!mkdir -p Drive/MyDrive/Armadillo-BaseOS-Guide/"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hpxWRUBnYRr8"
      },
      "source": [
        "## TensorBoard\n",
        "\n",
        "TensorFlowには、学習中の精度などのパラメータの推移を可視化するTensorBoardという機能があります。詳細は[公式ドキュメント](https://www.tensorflow.org/tensorboard?hl=ja)を参照してください。\n",
        "\n",
        "Google Colaboratory上でも使用可能ですので使ってみます。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XL2eAeqCnkaE"
      },
      "source": [
        "%load_ext tensorboard"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UY4GWbFmnpYE"
      },
      "source": [
        "tensorboard --logdir '/content/Drive/MyDrive/Armadillo-BaseOS-Guide'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3dUIK_WSZJDw"
      },
      "source": [
        "## 学習開始\n",
        "\n",
        "いよいよ転移学習を開始します。  \n",
        "ここまで準備してきたので、学習及び検証はObject Detection APIで用意されているPythonスクリプトに必要な情報を渡すだけで実行できます。\n",
        "\n",
        "大量のログが出力された後に、各ステップ毎の学習結果が出力され始めたら転移学習が正しく開始されたことになります。\n",
        "\n",
        "直前のセルで実行したTensorBoardで学習中の進捗が確認できます。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YP2DCH1Any_r"
      },
      "source": [
        "!python /content/models/research/object_detection/model_main_tf2.py \\\n",
        "    --pipeline_config_path=\"/content/working/03_pretrained_model/ssd_mobilenet_v2_320x320_coco17_tpu-8/pipeline.config\" \\\n",
        "    --model_dir=\"/content/Drive/MyDrive/Armadillo-BaseOS-Guide\" \\\n",
        "    --num_train_steps=100000 \\\n",
        "    --alsologtostderr \\\n",
        "    --sample_1_of_n_eval_examples=1 \\\n",
        "    --num_eval_steps=1000"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GVJH-ungatL4"
      },
      "source": [
        "# TFLite形式への変換\n",
        "\n",
        "ここまでの手順で作成した推移モデルは、Armadilloのようなエッジデバイス向けのモデルではありません。\n",
        "\n",
        "以下のスクリプトを実行することで、エッジデバイス向けに最適化されたTFLite形式の推移モデルに変換することができます。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q4HbfnNJcooP"
      },
      "source": [
        "## 変換の前処理\n",
        "\n",
        "転移学習で得た推移モデルは直接TFLite形式に変換することはできず、まずはTFLite形式への変換に適したSavedModel形式に変換します。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "L4AU08bqEcTh"
      },
      "source": [
        "%%bash\n",
        "python models/research/object_detection/export_tflite_graph_tf2.py \\\n",
        "  --pipeline_config_path \"/content/working/03_pretrained_model/ssd_mobilenet_v2_320x320_coco17_tpu-8/pipeline.config\" \\\n",
        "  --trained_checkpoint_dir \"/content/Drive/MyDrive/Armadillo-BaseOS-Guide\" \\\n",
        "  --output_directory \"/content/Drive/MyDrive/Armadillo-BaseOS-Guide/tflite\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xykCnuuscaTl"
      },
      "source": [
        "## 変換の実行\n",
        "\n",
        "前処理が完了したら、TFLite形式への変換を行います。\n",
        "\n",
        "ここでは単純に変換するのではなく、Armadillo-IoT G4搭載のNPU向けに整数量子化を行います。\n",
        "量子化を行うことでわずかにモデルの精度が低下しますが、モデルのサイズの削減や推論速度の向上、メモリ使用量の削減などが期待できます。\n",
        "\n",
        "モデルの量子化について、詳細は[公式ドキュメント](https://www.tensorflow.org/lite/performance/post_training_quantization?hl=ja)を合わせて参照してください。\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qu73ikEoQU8a"
      },
      "source": [
        "def representative_dataset():\n",
        "  dataset_list = tf.data.Dataset.list_files('/content/working/01_train_data/*.jpg')\n",
        "  for i in range(100): # ここの値は用意した画像数によって上下する(100~500程度が良いとされている)\n",
        "    image = next(iter(dataset_list))\n",
        "    image = tf.io.read_file(image)\n",
        "    image = tf.io.decode_jpeg(image, channels=3)\n",
        "    image = tf.image.resize(image, [300, 300])\n",
        "    image = tf.cast(image / 255, tf.float32)\n",
        "    image = tf.expand_dims(image, 0)\n",
        "    yield [image]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UwX9b23ZTQao"
      },
      "source": [
        "converter = tf.lite.TFLiteConverter.from_saved_model('/content/Drive/MyDrive/Armadillo-BaseOS-Guide/tflite/saved_model')\n",
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "converter.representative_dataset = representative_dataset\n",
        "converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]\n",
        "converter.inference_input_type = tf.uint8\n",
        "converter.inference_output_type = tf.uint8\n",
        "tflite_quant_model = converter.convert()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C6GvDSvoibpl"
      },
      "source": [
        "## 変換後の推論モデルの保存\n",
        "\n",
        "model.tfliteというファイル名で保存します。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sCVY4fHmS57E"
      },
      "source": [
        "open('/content/Drive/MyDrive/Armadillo-BaseOS-Guide/tflite/model.tflite', 'wb').write(tflite_quant_model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_aMl7L_5i7YX"
      },
      "source": [
        "# 生成したtfliteモデルの動作確認\n",
        "\n",
        "保存したmodel.tfliteを読み出し、実際に物体検出を行ってみます。\n",
        "\n",
        "以下の手順は、最終的なArmadillo-IoT ゲートウェイ G4上でmodel.tfliteを使用したアプリケーションを作成する際の推論モデルの取り扱いの参考としても使用できます。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JPGPUsDVjNMl"
      },
      "source": [
        "## model.tfliteの読み出し"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import tensorflow as tf\n",
        "import cv2\n",
        "import numpy as np\n",
        "import copy\n",
        "from google.colab.patches import cv2_imshow"
      ],
      "metadata": {
        "id": "Hlkdr1j2hQWf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "omHVbXUkUSB-"
      },
      "source": [
        "interpreter = tf.lite.Interpreter(model_path='/content/Drive/MyDrive/Armadillo-BaseOS-Guide/tflite/model.tflite')\n",
        "interpreter.allocate_tensors()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hlmWrl46jwAU"
      },
      "source": [
        "## TFLite形式のモデルから入出力情報を読み取る"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GhrJy0RxUgiH"
      },
      "source": [
        "# get_[input|output]_detailsで推論モデルの入出力情報を得ることができる\n",
        "input_details = interpreter.get_input_details()\n",
        "output_details = interpreter.get_output_details()\n",
        "\n",
        "# 入力画像サイズを取得\n",
        "_, input_height, input_width, _ = input_details[0]['shape']"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qm-xPoh8j_AP"
      },
      "source": [
        "## 推論"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LIkzu4bhj8CT"
      },
      "source": [
        "# 適当な画像をリサイズして入力\n",
        "imagefile = '/content/working/01_train_data/000001.jpg'\n",
        "input_data = cv2.imread(imagefile)\n",
        "resized = cv2.resize(input_data, dsize=(input_width, input_height))\n",
        "resized = resized / 255\n",
        "if input_details[0]['dtype'] == np.uint8:\n",
        "    i_scale, i_zeropoint = input_details[0]['quantization']\n",
        "    resized = (resized / i_scale) + i_zeropoint\n",
        "resized = np.expand_dims(resized, 0)\n",
        "interpreter.set_tensor(input_details[0]['index'], resized.astype(np.uint8))\n",
        "\n",
        "# 推論の実行\n",
        "interpreter.invoke()\n",
        "\n",
        "def get_output(index):\n",
        "    output = interpreter.get_tensor(output_details[index]['index'])[0]\n",
        "    if output_details[index]['dtype'] == np.uint8:\n",
        "        o_scale, o_zeropoint = output_details[index]['quantization']\n",
        "        output = (output - o_zeropoint) * o_scale\n",
        "    return output\n",
        "\n",
        "# 推論結果を取得して表示\n",
        "scores = get_output(0)\n",
        "boxes = get_output(1)\n",
        "\n",
        "for i in range(4):\n",
        "    print(\"{}: {}\".format(i, get_output(i)))\n",
        "\n",
        "# print(scores)\n",
        "# print(boxes)\n",
        "\n",
        "# 推論結果にしきい値を設定して抽出・整形\n",
        "threshold = 0.8\n",
        "results = []\n",
        "for i in range(len(scores)):\n",
        "  if scores[i] >= threshold:\n",
        "    res = {}\n",
        "    res['score'] = scores[i]\n",
        "    res['box'] = boxes[i]\n",
        "    results.append(res)\n",
        "# print(results)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UjyCqr0Akwwd"
      },
      "source": [
        "## 推論結果の描画"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xZKTBgA8jVPV"
      },
      "source": [
        "# 元の画像をOpenCVで読み出し\n",
        "image = cv2.imread(imagefile, cv2.IMREAD_UNCHANGED)\n",
        "image_height, image_width, _ = image.shape\n",
        "cp = copy.deepcopy(image)\n",
        "\n",
        "# 推論結果を整形して元画像の上に描画\n",
        "for result in results:\n",
        "  y1, x1, y2, x2 = result['box']\n",
        "  x1 = int(x1 * image_width)\n",
        "  x2 = int(x2 * image_width)\n",
        "  y1 = int(y1 * image_height)\n",
        "  y2 = int(y2 * image_height)\n",
        "  score = result['score']\n",
        "  cv2.rectangle(cp, (x1, y1), (x2, y2), (255, 0, 0), 2)\n",
        "  cv2.putText(cp, str('{:.2f}'.format(score)), (x1, y1-10), cv2.FONT_HERSHEY_PLAIN, 1.5, (255, 0, 0), 2, cv2.LINE_AA)\n",
        "\n",
        "cv2_imshow(cp)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j8uYZ9filGAV"
      },
      "source": [
        "# 推論モデルのダウンロード\n",
        "\n",
        "以下を実行することで、model.tfliteをローカルにダウンロードできます。\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2_9sdzj8DDFQ"
      },
      "source": [
        "from google.colab import files\n",
        "files.download('/content/Drive/MyDrive/Armadillo-BaseOS-Guide/tflite/model.tflite') "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "以上で推論モデルの生成は完了です。\n",
        "\n",
        "Armadillo-IoT ゲートウェイ G4への組み込み、アプリケーション開発については、引き続き「Armadillo Base OS 開発ガイド」を参照してください。"
      ],
      "metadata": {
        "id": "Bo4XExEphbYf"
      }
    }
  ]
}