Defcon4_TFLite_image_classification.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import absolute_import, division\n",
    "from IPython.display import clear_output\n",
    "from BashColors import C\n",
    "\n",
    "import glob, os, shutil\n",
    "from os.path import *\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '4'\n",
    "\n",
    "contentPath=os.getcwd()\n",
    "\n",
    "import numpy as np\n",
    "np.set_printoptions(precision=1)\n",
    "import tensorflow as tf\n",
    "assert tf.__version__.startswith('2')\n",
    "import tflite_model_maker\n",
    "from tflite_model_maker import model_spec\n",
    "from tflite_model_maker import image_classifier\n",
    "from tflite_model_maker.config import ExportFormat\n",
    "from tflite_model_maker.config import QuantizationConfig\n",
    "from tflite_model_maker.image_classifier import DataLoader\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "print(f'tf version: {C.BIBlue}{tf.__version__}{C.ColorOff}')\n",
    "print(f'tflite version: {C.BIBlue}{tflite_model_maker.__version__}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "from time import sleep\n",
    "\n",
    "def resizeImage(thisPath):\n",
    "    print(f'{C.BIPurple}{thisPath}{C.ColorOff}')\n",
    "    img = cv2.imread(thisPath)\n",
    "    img = tf.image.resize(\n",
    "        img, size=(224,224), method=tf.image.ResizeMethod.BILINEAR,\n",
    "        preserve_aspect_ratio=False,\n",
    "        antialias=False, name=None\n",
    "    )\n",
    "    img=np.copy(img)\n",
    "    cv2.imwrite(thisPath, img)\n",
    "    sleep(0.2)\n",
    "    print(img.shape)\n",
    "    \n",
    "os.chdir(contentPath)\n",
    "imageList=[]\n",
    "initialList=[]\n",
    "globList=[]\n",
    "\n",
    "globList=glob.glob('**', recursive=True)\n",
    "for pth in globList:\n",
    "    fullPath=abspath(pth)\n",
    "    # print(fullPath)\n",
    "    initialList.append(fullPath)\n",
    "    if fullPath.endswith('.png'):\n",
    "        imageList.append(fullPath)\n",
    "        \n",
    "for item in imageList:\n",
    "    # resizeImage(item)\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "I_fOlZsklmlL"
   },
   "outputs": [],
   "source": [
    "# Split it to training data (80%), validation data (10%, optional) and testing data (10%).\n",
    "\n",
    "data = DataLoader.from_folder('DataGenerator')\n",
    "print(f'data files: {C.BIPurple}{len(data)}{C.ColorOff}')\n",
    "train_data, validation_data = data.split(0.8)\n",
    "print(f'train_data: {C.BIPurple}{len(train_data)}{C.ColorOff}')\n",
    "print(f'validation_data: {C.BIPurple}{len(validation_data)}{C.ColorOff}\\n')\n",
    "\n",
    "testData = DataLoader.from_folder('images')\n",
    "test_data = testData\n",
    "print(f'test_data: {C.BIPurple}{len(test_data)}{C.ColorOff}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ih4Wx44I482b"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12, 12))\n",
    "for i, (image, label) in enumerate(train_data.gen_dataset().unbatch().take(3)):\n",
    "  plt.subplot(3, 3, i+1)\n",
    "  plt.xticks([])\n",
    "  plt.yticks([])\n",
    "  plt.grid(False)\n",
    "  plt.imshow(image.numpy(), cmap=plt.cm.gray)\n",
    "  plt.xlabel(data.index_to_label[label.numpy()])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "epochs: more epochs could achieve better accuracy until it converges but training for too many epochs may lead to overfitting.\n",
    "    \n",
    "dropout_rate: The rate for dropout, avoid overfitting. None by default.\n",
    "\n",
    "batch_size: number of samples to use in one training step. None by default.\n",
    "    \n",
    "validation_data: Validation data. If None, skips validation process. None by default.\n",
    "    \n",
    "train_whole_model: If true, the Hub module is trained together with the classification layer on top. Otherwise, only train the top classification layer. None by default.\n",
    "    \n",
    "learning_rate: Base learning rate. None by default.\n",
    "    \n",
    "momentum: a Python float forwarded to the optimizer. Only used when use_hub_library is True. None by default.\n",
    "    \n",
    "shuffle: Boolean, whether the data should be shuffled. False by default.\n",
    "    \n",
    "use_augmentation: Boolean, use data augmentation for preprocessing. False by default.\n",
    "    \n",
    "use_hub_library: Boolean, use make_image_classifier_lib from tensorflow hub to retrain the model. This training pipeline could achieve better performance for complicated dataset with many categories. True by default.\n",
    "    \n",
    "warmup_steps: Number of warmup steps for warmup schedule on learning rate. If None, the default warmup_steps is used which is the total training steps in two epochs. Only used when use_hub_library is False. None by default.\n",
    "    \n",
    "model_dir: Optional, the location of the model checkpoint files. Only used when use_hub_library is False. None by default."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE =8\n",
    "EPOCHS = 10\n",
    "DROPOUT_RATE = 0.5\n",
    "LEARNING_RATE = 0.001\n",
    "\n",
    "model = image_classifier.create(\n",
    "    train_data = train_data,\n",
    "    validation_data = validation_data,\n",
    "    model_spec = model_spec.get('mobilenet_v2'),\n",
    "    learning_rate = LEARNING_RATE,\n",
    "    batch_size = BATCH_SIZE,\n",
    "    epochs = EPOCHS,\n",
    "    dropout_rate = DROPOUT_RATE,\n",
    "    use_augmentation = True, # default=False\n",
    "    use_hub_library = True,  # default=True\n",
    "    momentum = 0.9, # Only used when use_hub_library is True\n",
    "    shuffle = True,\n",
    "    train_whole_model = False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LP5FPk_tOxoZ"
   },
   "source": [
    "### Step 3: Evaluate the Customized Model\n",
    "\n",
    "Evaluate the result of the model, get the loss and accuracy of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "A8c2ZQ0J3Riy"
   },
   "outputs": [],
   "source": [
    "loss, accuracy = model.evaluate(test_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6ZCrYOWoCt05"
   },
   "source": [
    "We could plot the predicted results in 100 test images. Predicted labels with red color are the wrong predicted results while others are correct."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aeHoGAceO2xV"
   },
   "source": [
    "### Step 4: Export to TensorFlow Lite Model\n",
    "\n",
    "Convert the trained model to TensorFlow Lite model format with [metadata](https://www.tensorflow.org/lite/convert/metadata) so that you can later use in an on-device ML application. The label file and the vocab file are embedded in metadata. The default TFLite filename is `model.tflite`.\n",
    "\n",
    "In many on-device ML application, the model size is an important factor. Therefore, it is recommended that you apply quantize the model to make it smaller and potentially run faster.\n",
    "The default post-training quantization technique is full integer quantization for the image classification task."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ROS2Ay2jMPCl"
   },
   "source": [
    "See [example applications and guides of image classification](https://www.tensorflow.org/lite/models/image_classification/overview#example_applications_and_guides) for more details about how to integrate the TensorFlow Lite model into mobile apps.\n",
    "\n",
    "This model can be integrated into an Android or an iOS app using the [ImageClassifier API](https://www.tensorflow.org/lite/inference_with_metadata/task_library/image_classifier) of the [TensorFlow Lite Task Library](https://www.tensorflow.org/lite/inference_with_metadata/task_library/overview)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Im6wA9lK3TQB"
   },
   "outputs": [],
   "source": [
    "# model.export(export_dir='.')\n",
    "\n",
    "exportList = [ExportFormat.TFLITE,\n",
    "              ExportFormat.LABEL,\n",
    "              ExportFormat.SAVED_MODEL]\n",
    "float16Config = QuantizationConfig.for_float16()\n",
    "# model.export(export_dir='.', export_format = ExportFormat.LABEL)\n",
    "model.export(export_dir = contentPath,\n",
    "             tflite_filename = 'Defcon4_fp16.tflite',\n",
    "             export_format = exportList,\n",
    "             config=float16Config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "S1YoPX5wOK-u"
   },
   "outputs": [],
   "source": [
    "model.evaluate_tflite('Defcon4_fp16.tflite', test_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "n9O9Kx7nDQWD"
   },
   "outputs": [],
   "source": [
    "# A helper function that returns 'red'/'black' depending on if its two input\n",
    "# parameter matches or not.\n",
    "def get_label_color(val1, val2):\n",
    "  if val1 == val2:\n",
    "    return 'blue'\n",
    "  else:\n",
    "    return 'black'\n",
    "\n",
    "# Then plot 100 test images and their predicted labels.\n",
    "# If a prediction result is different from the label provided label in \"test\"\n",
    "# dataset, we will highlight it in red color.\n",
    "plt.figure(figsize=(224/4, 224/4))\n",
    "predicts = model.predict_top_k(test_data)\n",
    "for i, (image, label) in enumerate(test_data.gen_dataset().unbatch().take(100)):\n",
    "    ax = plt.subplot(11, 11, i+1)\n",
    "    plt.xticks([])\n",
    "    plt.yticks([])\n",
    "    plt.grid(False)\n",
    "    plt.imshow(image.numpy(), cmap=plt.cm.gray)\n",
    "    predict_label = predicts[i][0][0]\n",
    "    color = get_label_color(predict_label,\n",
    "                            test_data.index_to_label[label.numpy()])\n",
    "    ax.xaxis.label.set_fontsize(24)\n",
    "    ax.xaxis.label.set_color(color)\n",
    "    plt.xlabel('%s' % predict_label)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(os.getcwd())\n",
    "import glob\n",
    "newList=[]\n",
    "globList=[]\n",
    "globList=glob.glob('**', recursive=True)\n",
    "\n",
    "for pth in globList:\n",
    "    fullPath=abspath(pth)\n",
    "    if not fullPath in initialList:\n",
    "        newList.append(fullPath)\n",
    "        if isdir(fullPath):\n",
    "            print(f'{C.BIBlue}{fullPath}{C.ColorOff}')\n",
    "        elif isfile(fullPath):\n",
    "            print(fullPath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# q"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zNDBP2qA54aK"
   },
   "source": [
    "## Advanced Usage\n",
    "\n",
    "The `create` function is the critical part of this library. It uses transfer learning with a pretrained model similar to the [tutorial](https://www.tensorflow.org/tutorials/images/transfer_learning).\n",
    "\n",
    "The `create` function contains the following steps:\n",
    "\n",
    "1.   Split the data into training, validation, testing data according to parameter `validation_ratio` and `test_ratio`. The default value of `validation_ratio` and `test_ratio` are `0.1` and `0.1`.\n",
    "2.   Download a [Image Feature Vector](https://www.tensorflow.org/hub/common_signatures/images#image_feature_vector) as the base model from TensorFlow Hub. The default pre-trained model is  EfficientNet-Lite0.\n",
    "3.   Add a classifier head with a Dropout Layer with `dropout_rate` between head layer and pre-trained model. The default `dropout_rate` is the default `dropout_rate` value from [make_image_classifier_lib](https://github.com/tensorflow/hub/blob/master/tensorflow_hub/tools/make_image_classifier/make_image_classifier_lib.py#L55) by TensorFlow Hub.\n",
    "4.   Preprocess the raw input data. Currently, preprocessing steps including normalizing the value of each image pixel to model input scale and resizing it to model input size.   EfficientNet-Lite0 have the input scale `[0, 1]` and the input image size `[224, 224, 3]`.\n",
    "5.   Feed the data into the classifier model. By default, the training parameters such as training epochs, batch size, learning rate, momentum are the default values from [make_image_classifier_lib](https://github.com/tensorflow/hub/blob/master/tensorflow_hub/tools/make_image_classifier/make_image_classifier_lib.py#L55) by TensorFlow Hub. Only the classifier head is trained.\n",
    "\n",
    "\n",
    "In this section, we describe several advanced topics, including switching to a different image classification model, changing the training hyperparameters etc.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Gc4Jk8TvBQfm"
   },
   "source": [
    "## Customize Post-training quantization on the TensorFLow Lite model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tD8BOYrHBiDt"
   },
   "source": [
    "[Post-training quantization](https://www.tensorflow.org/lite/performance/post_training_quantization) is a conversion technique that can reduce model size and inference latency, while also improving CPU and hardware accelerator inference speed, with a little degradation in model accuracy. Thus, it's widely used to optimize the model.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iyIo0d5TCzE2"
   },
   "source": [
    "Model Maker library applies a default post-training quantization techique when exporting the model. If you want to customize post-training quantization, Model Maker supports multiple post-training quantization options using [QuantizationConfig](https://www.tensorflow.org/lite/api_docs/python/tflite_model_maker/config/QuantizationConfig) as well. Let's take float16 quantization as an instance. First, define the quantization config."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WTJzFQnJFMjr"
   },
   "outputs": [],
   "source": [
    "config = QuantizationConfig.for_float16()\n",
    "\n",
    "model.export(export_dir =contentPath,\n",
    "             tflite_filename = 'Defcon4_fp16.tflite',\n",
    "             quantization_config=config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lB2Go3HW8X7_"
   },
   "outputs": [],
   "source": [
    "loss, accuracy = model.evaluate(test_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UhZ5IRKdeex3"
   },
   "source": [
    "### Change your own custom model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "svTjlZhrCrcV"
   },
   "source": [
    "If we'd like to use the custom model that's not in TensorFlow Hub, we should create and export [ModelSpec](https://www.tensorflow.org/hub/api_docs/python/hub/ModuleSpec) in TensorFlow Hub.\n",
    "\n",
    "Then start to define `ModelSpec` object like the process above."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dhBU5NCy5Ji2"
   },
   "source": [
    "# Read more\n",
    "\n",
    "You can read our [image classification](https://www.tensorflow.org/lite/examples/image_classification/overview) example to learn technical details. For more information, please refer to:\n",
    "\n",
    "*   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",
    "*   Task Library: [ImageClassifier](https://www.tensorflow.org/lite/inference_with_metadata/task_library/image_classifier) for deployment.\n",
    "*   The end-to-end reference apps: [Android](https://github.com/tensorflow/examples/tree/master/lite/examples/image_classification/android), [iOS](https://github.com/tensorflow/examples/tree/master/lite/examples/image_classification/ios), and [Raspberry PI](https://github.com/tensorflow/examples/tree/master/lite/examples/image_classification/raspberry_pi).\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Model Maker Image Classification Tutorial",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}