4-8-11-CustomMagicWand.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "4-8-11-CustomMagicWand",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rK0PP6cDJQEC"
      },
      "source": [
        "# Training Your Custom Magic Wand Model\r\n",
        "\r\n",
        "It is now time for you to train your own custom magic wand model using your Custom Dataset!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V4asiPfAJaQW"
      },
      "source": [
        "##Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7aV25Y-uMBFc"
      },
      "source": [
        "### Import Packages and Set Constants"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yHO0QoaKInIw"
      },
      "source": [
        "%tensorflow_version 2.x\r\n",
        "import tensorflow as tf\r\n",
        "from tensorflow import keras\r\n",
        "from tensorflow.keras import layers\r\n",
        "from tensorflow.keras.preprocessing import image_dataset_from_directory\r\n",
        "from google.colab import files\r\n",
        "from IPython.display import Image, display\r\n",
        "import numpy as np\r\n",
        "import pandas as pd\r\n",
        "import matplotlib.pyplot as plt\r\n",
        "from pathlib import Path\r\n",
        "import shutil\r\n",
        "import PIL\r\n",
        "import math\r\n",
        "import glob\r\n",
        "import json\r\n",
        "import os\r\n",
        "!apt-get update && apt-get -qq install xxd"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0g1pF6RfViPr"
      },
      "source": [
        "# Define filenames and set up directory structure\n",
        "MODELS_DIR = 'models'\n",
        "if not os.path.exists(MODELS_DIR):\n",
        "  os.mkdir(MODELS_DIR)\n",
        "SAVED_MODEL_FILENAME = os.path.join(MODELS_DIR, \"magic_wand\")\n",
        "FLOAT_TFL_MODEL_FILENAME = os.path.join(MODELS_DIR, \"magic_wand_float.tfl\")\n",
        "QUANTIZED_TFL_MODEL_FILENAME = os.path.join(MODELS_DIR, \"magic_wand.tfl\")\n",
        "TFL_CC_MODEL_FILENAME = os.path.join(MODELS_DIR, \"magic_wand.cc\")\n",
        "\n",
        "DATASET_DIR =  'dataset'\n",
        "if not os.path.exists(DATASET_DIR):\n",
        "  os.mkdir(DATASET_DIR)\n",
        "TRAIN_DIR = os.path.join(DATASET_DIR, \"train\")\n",
        "VAL_DIR = os.path.join(DATASET_DIR, \"validation\")\n",
        "TEST_DIR = os.path.join(DATASET_DIR, \"test\")\n",
        "!rm -rf sample_data\n",
        "\n",
        "CHKPT_DIR =  'checkpoints'\n",
        "if not os.path.exists(CHKPT_DIR):\n",
        "  os.mkdir(CHKPT_DIR)\n",
        "\n",
        "# Train Split\n",
        "TEST_PERCENTAGE = 10\n",
        "VALIDATION_PERCENTAGE = 10\n",
        "TRAIN_PERCENTAGE = 100 - (TEST_PERCENTAGE + VALIDATION_PERCENTAGE)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rz4y_wQnKb11"
      },
      "source": [
        "### Load Your Custom Dataset\r\n",
        "Now you'll need to upload all of your custom gesture files that you created using the Magic Wand tool (aka the ```*.json``` files). **Note: you can select multiple files and upload them all at once!** \r\n",
        "\r\n",
        "If you are having trouble uploading files because your internet bandwidth is too slow feel free to uncomment the lines below to instead use Pete's digits dataset."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xvPo0OP0TFDq"
      },
      "source": [
        "# Upload your files\n",
        "os.chdir(\"/content/dataset\")\n",
        "uploaded = files.upload()\n",
        "os.chdir(\"/content\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z_4rbb0WSSFV"
      },
      "source": [
        "# Or use Pete's Digits Dataset\r\n",
        "# !curl -L https://github.com/petewarden/magic_wand_digit_data/archive/8170591863f9addca27b1a963263f7c7bed33f41.zip -o magic_wand_digit_data.zip\r\n",
        "# !unzip magic_wand_digit_data.zip\r\n",
        "# !rm -rf magic_wand_digit_data.zip\r\n",
        "# !mv magic_wand_digit_data-*/* dataset\r\n",
        "# !rm -rf magic_wand_digit_data-*"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R4loOL1hPM_2"
      },
      "source": [
        "**Update the variable below with the number of labeled gestures in your dataset** \r\n",
        "Note: Use the number of unique gestures/labels and *not* the number of samples in your dataset."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9miJi7WGPNft"
      },
      "source": [
        "NUM_GESTURES = 10 # UPDATE THIS WITH THE NUMBER OF UNIQUE GESTURES IN YOUR DATASET #"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kJBM70APMENw"
      },
      "source": [
        "Next we'll parse the JSON files into a python object which we can more easily work with."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mWO58-igVFSd"
      },
      "source": [
        "dataset_jsons = DATASET_DIR + \"/*.json\"\n",
        "strokes = []\n",
        "for filename in glob.glob(dataset_jsons):\n",
        "  with open(filename, \"r\") as file:\n",
        "    file_contents = file.read()\n",
        "  file_data = json.loads(file_contents)\n",
        "  for stroke in file_data[\"strokes\"]:\n",
        "    stroke[\"filename\"] = filename\n",
        "    strokes.append(stroke)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NA8KOrP0MTMt"
      },
      "source": [
        "If you'd like to visualize any of your gestures you can use the helper function below!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xfLzrpyLVJ5S"
      },
      "source": [
        "# Helper function to visualize the data\n",
        "def plot_stroke(stroke):\n",
        "  x_array = []\n",
        "  y_array = []\n",
        "  for coords in stroke[\"strokePoints\"]:\n",
        "    x_array.append(coords[\"x\"])\n",
        "    y_array.append(coords[\"y\"])\n",
        "\n",
        "  fig = plt.figure(figsize=(12.8, 4.8))\n",
        "  fig.suptitle(stroke[\"label\"])\n",
        "\n",
        "  ax = fig.add_subplot(131)\n",
        "  ax.set_xlabel('x')\n",
        "  ax.set_ylabel('y')\n",
        "  ax.set_xlim(-0.4, 0.4)\n",
        "  ax.set_ylim(-0.4, 0.4)\n",
        "  ax.plot(x_array, y_array)\n",
        "\n",
        "  plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dveZd2ZuW-jl"
      },
      "source": [
        "# Display a stroke from the strokes python variable\r\n",
        "plot_stroke(strokes[0])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ml1UYg-oMpQo"
      },
      "source": [
        "### Preprocess your Dataset\r\n",
        "Next we'll preprocess the dataset to prepare it for training. By preprocessing the data in bulk before training the whole training process will execute much faster. To do so, we'll convert the strokes into rastered images using the helper functions below. This is the process used in real-time in the Arduino code to convert a gesture into an image that the CNN we are going to train can then process.\r\n",
        "\r\n",
        "Once we have converted the dataset to rasterized images we will generate a ```Keras``` dataset for use in training."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3FVPj-eqjvoB"
      },
      "source": [
        "FIXED_POINT = 256\n",
        "\n",
        "def mul_fp(a, b):\n",
        "  return (a * b) / FIXED_POINT\n",
        "\n",
        "def div_fp(a, b):\n",
        "  if b == 0:\n",
        "    b = 1\n",
        "  return (a * FIXED_POINT) / b\n",
        "\n",
        "def float_to_fp(a):\n",
        "  return math.floor(a * FIXED_POINT)\n",
        "\n",
        "def norm_to_coord_fp(a, range_fp, half_size_fp):\n",
        "  a_fp = float_to_fp(a)\n",
        "  norm_fp = div_fp(a_fp, range_fp)\n",
        "  return mul_fp(norm_fp, half_size_fp) + half_size_fp\n",
        "\n",
        "def round_fp_to_int(a):\n",
        "  return math.floor((a + (FIXED_POINT / 2)) / FIXED_POINT)\n",
        "\n",
        "def gate(a, min, max):\n",
        "  if a < min:\n",
        "    return min\n",
        "  elif a > max:\n",
        "    return max\n",
        "  else:\n",
        "    return a\n",
        "\n",
        "def rasterize_stroke(stroke_points, x_range, y_range, width, height):\n",
        "  num_channels = 3\n",
        "  buffer_byte_count = height * width * num_channels\n",
        "  buffer = bytearray(buffer_byte_count)\n",
        "\n",
        "  width_fp = width * FIXED_POINT\n",
        "  height_fp = height * FIXED_POINT\n",
        "  half_width_fp = width_fp / 2\n",
        "  half_height_fp = height_fp / 2\n",
        "  x_range_fp = float_to_fp(x_range)\n",
        "  y_range_fp = float_to_fp(y_range)\n",
        "\n",
        "  t_inc_fp = FIXED_POINT / len(stroke_points)\n",
        "\n",
        "  one_half_fp = (FIXED_POINT / 2)\n",
        "\n",
        "  for point_index in range(len(stroke_points) - 1):\n",
        "    start_point = stroke_points[point_index]\n",
        "    end_point = stroke_points[point_index + 1]\n",
        "    start_x_fp = norm_to_coord_fp(start_point[\"x\"], x_range_fp, half_width_fp)\n",
        "    start_y_fp = norm_to_coord_fp(-start_point[\"y\"], y_range_fp, half_height_fp)\n",
        "    end_x_fp = norm_to_coord_fp(end_point[\"x\"], x_range_fp, half_width_fp)\n",
        "    end_y_fp = norm_to_coord_fp(-end_point[\"y\"], y_range_fp, half_height_fp)\n",
        "    delta_x_fp = end_x_fp - start_x_fp\n",
        "    delta_y_fp = end_y_fp - start_y_fp\n",
        "\n",
        "    t_fp = point_index * t_inc_fp\n",
        "    if t_fp < one_half_fp:\n",
        "      local_t_fp = div_fp(t_fp, one_half_fp)\n",
        "      one_minus_t_fp = FIXED_POINT - local_t_fp\n",
        "      red = round_fp_to_int(one_minus_t_fp * 255)\n",
        "      green = round_fp_to_int(local_t_fp * 255)\n",
        "      blue = 0\n",
        "    else:\n",
        "      local_t_fp = div_fp(t_fp - one_half_fp, one_half_fp)\n",
        "      one_minus_t_fp = FIXED_POINT - local_t_fp\n",
        "      red = 0\n",
        "      green = round_fp_to_int(one_minus_t_fp * 255)\n",
        "      blue = round_fp_to_int(local_t_fp * 255)\n",
        "    red = gate(red, 0, 255)\n",
        "    green = gate(green, 0, 255)\n",
        "    blue = gate(blue, 0, 255)\n",
        "\n",
        "    if abs(delta_x_fp) > abs(delta_y_fp):\n",
        "      line_length = abs(round_fp_to_int(delta_x_fp))\n",
        "      if delta_x_fp > 0:\n",
        "        x_inc_fp = 1 * FIXED_POINT\n",
        "        y_inc_fp = div_fp(delta_y_fp, delta_x_fp)\n",
        "      else:\n",
        "        x_inc_fp = -1 * FIXED_POINT\n",
        "        y_inc_fp = -div_fp(delta_y_fp, delta_x_fp)\n",
        "    else:\n",
        "      line_length = abs(round_fp_to_int(delta_y_fp))\n",
        "      if delta_y_fp > 0:\n",
        "        y_inc_fp = 1 * FIXED_POINT\n",
        "        x_inc_fp = div_fp(delta_x_fp, delta_y_fp)\n",
        "      else:\n",
        "        y_inc_fp = -1 * FIXED_POINT\n",
        "        x_inc_fp = -div_fp(delta_x_fp, delta_y_fp)\n",
        "    for i in range(line_length + 1):\n",
        "      x_fp = start_x_fp + (i * x_inc_fp)\n",
        "      y_fp = start_y_fp + (i * y_inc_fp)\n",
        "      x = round_fp_to_int(x_fp)\n",
        "      y = round_fp_to_int(y_fp)\n",
        "      if (x < 0) or (x >= width) or (y < 0) or (y >= height):\n",
        "        continue\n",
        "      buffer_index = (y * width * num_channels) + (x * num_channels)\n",
        "      buffer[buffer_index + 0] = red\n",
        "      buffer[buffer_index + 1] = green\n",
        "      buffer[buffer_index + 2] = blue\n",
        "  \n",
        "  np_buffer = np.frombuffer(buffer, dtype=np.uint8).reshape(height, width, num_channels)\n",
        "\n",
        "  return np_buffer"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "o-FOVdFpgkdf"
      },
      "source": [
        "X_RANGE = 0.6\n",
        "Y_RANGE = 0.6\n",
        "\n",
        "def ensure_empty_dir(dirname):\n",
        "  dirpath = Path(dirname)\n",
        "  if dirpath.exists() and dirpath.is_dir():\n",
        "    shutil.rmtree(dirpath)\n",
        "  dirpath.mkdir()\n",
        "\n",
        "def augment_points(points, move_range, scale_range, rotate_range):\n",
        "  move_x = np.random.uniform(low=-move_range, high=move_range)\n",
        "  move_y = np.random.uniform(low=-move_range, high=move_range)\n",
        "  scale = np.random.uniform(low=1.0-scale_range, high=1.0+scale_range)\n",
        "  rotate = np.random.uniform(low=-rotate_range, high=rotate_range)\n",
        "\n",
        "  x_axis_x = math.cos(rotate) * scale\n",
        "  x_axis_y = math.sin(rotate) * scale\n",
        "\n",
        "  y_axis_x = -math.sin(rotate) * scale\n",
        "  y_axis_y = math.cos(rotate) * scale\n",
        "\n",
        "  new_points = []\n",
        "  for point in points:\n",
        "    old_x = point[\"x\"]\n",
        "    old_y = point[\"y\"]\n",
        "    new_x = (x_axis_x * old_x) + (x_axis_y * old_y) + move_x\n",
        "    new_y = (y_axis_x * old_x) + (y_axis_y * old_y) + move_y\n",
        "    new_points.append({\"x\": new_x, \"y\": new_y})\n",
        "\n",
        "  return new_points\n",
        "\n",
        "def save_strokes_as_images(strokes, root_folder, width, height, augment_count):\n",
        "  ensure_empty_dir(root_folder)\n",
        "  labels = set()\n",
        "  for stroke in strokes:\n",
        "    labels.add(stroke[\"label\"].lower())\n",
        "  for label in labels:\n",
        "    label_path = Path(root_folder, label)\n",
        "    ensure_empty_dir(label_path)\n",
        "\n",
        "  label_counts = {}\n",
        "  for stroke in strokes:\n",
        "    points = stroke[\"strokePoints\"]\n",
        "    label = stroke[\"label\"].lower()\n",
        "    if label == \"\":\n",
        "      raise Exception(\"Missing label for %s:%d\" % (stroke[\"filename\"], stroke[\"index\"]))\n",
        "    if label not in label_counts:\n",
        "      label_counts[label] = 0\n",
        "    label_count = label_counts[label]\n",
        "    label_counts[label] += 1\n",
        "    raster = rasterize_stroke(points, X_RANGE, Y_RANGE, width, height)\n",
        "    image = PIL.Image.fromarray(raster)\n",
        "    image.save(Path(root_folder, label, str(label_count) + \".png\"))\n",
        "    for i in range(augment_count):\n",
        "      augmented_points = augment_points(points, 0.1, 0.1, 0.3)\n",
        "      raster = rasterize_stroke(augmented_points, X_RANGE, Y_RANGE, width, height)\n",
        "      image = PIL.Image.fromarray(raster)\n",
        "      image.save(Path(root_folder, label, str(label_count) + \"_a\" + str(i) + \".png\"))\n",
        "  return labels"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q7_LAzqkPpO8"
      },
      "source": [
        "Take the dataset and shuffle it into the Training/Validation/Test splits"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2cmmhBwW9d_p"
      },
      "source": [
        "IMAGE_WIDTH = 32\n",
        "IMAGE_HEIGHT = 32\n",
        "\n",
        "shuffled_strokes = strokes\n",
        "np.random.shuffle(shuffled_strokes)\n",
        "\n",
        "test_count = math.floor((len(shuffled_strokes) * TEST_PERCENTAGE) / 100)\n",
        "validation_count = math.floor((len(shuffled_strokes) * VALIDATION_PERCENTAGE) / 100)\n",
        "test_strokes = shuffled_strokes[0:test_count]\n",
        "validation_strokes = shuffled_strokes[test_count:(test_count + validation_count)]\n",
        "train_strokes = shuffled_strokes[(test_count + validation_count):]\n",
        "\n",
        "labels_test  = save_strokes_as_images(test_strokes, TEST_DIR, IMAGE_WIDTH, IMAGE_HEIGHT, 10)\n",
        "labels_val   = save_strokes_as_images(validation_strokes, VAL_DIR, IMAGE_WIDTH, IMAGE_HEIGHT, 0)\n",
        "labels_train = save_strokes_as_images(train_strokes, TRAIN_DIR, IMAGE_WIDTH, IMAGE_HEIGHT, 10)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VkARnseMRC4c"
      },
      "source": [
        "Also get the alphanumeric ordering of the labels as the Nueral Network will output its result in this order for the predicted class. **Make a note of this ordering as you will need to enter the labels in order in the Arduino code!**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JmVI7xD9QnqH"
      },
      "source": [
        "labels = sorted(labels_test.union(labels_val).union(labels_train))\n",
        "# get the conversion from label string to int\n",
        "labelToInt = {}\n",
        "currInt = 0\n",
        "for label in labels:\n",
        "  labelToInt[label] = currInt\n",
        "  currInt = currInt + 1\n",
        "intToLabel = {v: k for k, v in labelToInt.items()}\n",
        "print(intToLabel)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1YGHnw3xNlP_"
      },
      "source": [
        "If you'd like to visualize the difference between a stroke and its rasterized output image, run the cell below!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Rb6FnoKUNuvt"
      },
      "source": [
        "plot_stroke(strokes[0])\r\n",
        "raster = rasterize_stroke(strokes[0][\"strokePoints\"], 0.5, 0.5, 32, 32)\r\n",
        "PIL.Image.fromarray(raster).resize((512, 512), PIL.Image.NEAREST)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4WyU8c4KOBoj"
      },
      "source": [
        "Finally, we'll generate a dataset in ```Keras```."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9-Ttfz7LlPil"
      },
      "source": [
        "validation_ds = image_dataset_from_directory(\n",
        "    directory=VAL_DIR,\n",
        "    labels='inferred',\n",
        "    label_mode='categorical',\n",
        "    batch_size=32,\n",
        "    image_size=(IMAGE_WIDTH, IMAGE_HEIGHT)).prefetch(buffer_size=32)\n",
        "\n",
        "train_ds = image_dataset_from_directory(\n",
        "    directory=TRAIN_DIR,\n",
        "    labels='inferred',\n",
        "    label_mode='categorical',\n",
        "    batch_size=32,\n",
        "    image_size=(IMAGE_WIDTH, IMAGE_HEIGHT)).prefetch(buffer_size=32)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Q_vUMVmQn400"
      },
      "source": [
        "# Plot 9 of our final dataset items\n",
        "plt.figure(figsize=(10, 10))\n",
        "for images, labels in train_ds.take(1):\n",
        "  for i in range(9):\n",
        "    ax = plt.subplot(3, 3, i + 1)\n",
        "    plt.imshow(images[i].numpy().astype(\"uint8\"))\n",
        "    plt.axis(\"off\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fuQWd_spOwNl"
      },
      "source": [
        "## Define your Model\r\n",
        "\r\n",
        "Next we will define and visualize the CNN model that we will use for the magic wand!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "21qi3bLAo80t"
      },
      "source": [
        "def make_model(input_shape, num_classes):\n",
        "    inputs = keras.Input(shape=input_shape)\n",
        "\n",
        "    # Entry block\n",
        "    x = layers.experimental.preprocessing.Rescaling(1.0 / 255)(inputs)\n",
        "    x = layers.Conv2D(16, 3, strides=2, padding=\"same\")(x)\n",
        "    x = layers.BatchNormalization()(x)\n",
        "    x = layers.Activation(\"relu\")(x)\n",
        "    x = layers.Dropout(0.5)(x)\n",
        "\n",
        "    x = layers.Conv2D(32, 3, strides=2, padding=\"same\")(x)\n",
        "    x = layers.BatchNormalization()(x)\n",
        "    x = layers.Activation(\"relu\")(x)\n",
        "    x = layers.Dropout(0.5)(x)\n",
        "\n",
        "    x = layers.Conv2D(64, 3, strides=2, padding=\"same\")(x)\n",
        "    x = layers.BatchNormalization()(x)\n",
        "    x = layers.Activation(\"relu\")(x)\n",
        "    x = layers.Dropout(0.5)(x)\n",
        "\n",
        "    x = layers.GlobalAveragePooling2D()(x)\n",
        "    activation = \"softmax\"\n",
        "    units = num_classes\n",
        "\n",
        "    x = layers.Dropout(0.5)(x)\n",
        "    outputs = layers.Dense(units, activation=activation)(x)\n",
        "    return keras.Model(inputs, outputs)\n",
        "\n",
        "model = make_model(input_shape=(IMAGE_WIDTH, IMAGE_HEIGHT, 3), num_classes=NUM_GESTURES)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7DjpCkt0qZiy"
      },
      "source": [
        "# View the model layers as a diagram\n",
        "keras.utils.plot_model(model, show_shapes=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e4cdZX5JQapa"
      },
      "source": [
        "## Train your Model\r\n",
        "\r\n",
        "Now that we have a preprocessed dataset and a model its time to train that model with that dataset!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "287Am2coqkQC"
      },
      "source": [
        "# How many epochs to train for, we have found ~30 to be a good starting point\n",
        "EPOCHS = 30\n",
        "\n",
        "# Callback to save model checkpoints for future inspection or training\n",
        "checkpointFileLoc = CHKPT_DIR + \"/save_at_{epoch:02d}.h5\"\n",
        "modelCheckpointCallback = keras.callbacks.ModelCheckpoint(checkpointFileLoc)\n",
        "\n",
        "# Compile the model!\n",
        "model.compile(\n",
        "    optimizer=keras.optimizers.Adam(1e-3),\n",
        "    loss=\"binary_crossentropy\",\n",
        "    metrics=[\"accuracy\"],\n",
        ")\n",
        "\n",
        "# Run training\n",
        "history = model.fit(train_ds, epochs=EPOCHS, validation_data=validation_ds,\n",
        "                    callbacks=[modelCheckpointCallback])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "d26wGJn0t20g",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "071eb950-a8db-4dc3-b80b-717d22f19545"
      },
      "source": [
        "# save the model file\r\n",
        "model.save(SAVED_MODEL_FILENAME)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Assets written to: models/magic_wand/assets\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MvvyTSLhRYGJ"
      },
      "source": [
        "## Test your TensorFlow Model\r\n",
        "\r\n",
        "Lets now test out the TF model on the test dataset. We'll print out any gesture we get wrong as well as the percentage of known gestures correct as well as the number of gestures that were marked as unknown."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MYyLaqOYtxTH"
      },
      "source": [
        "SCORE_THRESHOLD = 0.75 # Confidence threshold to discard an image as \"unknown\"\n",
        "\n",
        "def predict_image(model, filename):\n",
        "  img = keras.preprocessing.image.load_img(filename, target_size=(IMAGE_WIDTH, IMAGE_HEIGHT))\n",
        "  img_array = keras.preprocessing.image.img_to_array(img)\n",
        "  img_array = tf.expand_dims(img_array, 0)  # Create batch axis\n",
        "  predictions = model.predict(img_array).flatten()\n",
        "  predicted_label_index = np.argmax(predictions)\n",
        "  predicted_score = predictions[predicted_label_index]\n",
        "  return (predicted_label_index, predicted_score)\n",
        "\n",
        "correct_count = 0\n",
        "wrong_count = 0\n",
        "discarded_count = 0\n",
        "for label_dir in glob.glob(TEST_DIR + \"/*\"):\n",
        "  label = label_dir.replace(TEST_DIR + \"/\", \"\")\n",
        "  print(\"Testing Gesture: \",label,\" with datasize: \",len(glob.glob(label_dir + \"/*.png\")))\n",
        "  for filename in glob.glob(label_dir + \"/*.png\"):\n",
        "    index, score = predict_image(model, filename)\n",
        "    if score < SCORE_THRESHOLD:\n",
        "      discarded_count += 1\n",
        "      continue\n",
        "    if index == labelToInt[label]:\n",
        "      correct_count += 1\n",
        "    else:\n",
        "      wrong_count += 1\n",
        "      print(label,index,score)\n",
        "      print(\"[%s] expected, [%s] found with score [%f]\" % (label, intToLabel[index], score))\n",
        "      display(Image(filename=filename))\n",
        "\n",
        "if correct_count + wrong_count == 0:\n",
        "  print(\"All images marked as unknown!\")\n",
        "else:\n",
        "  correct_percentage = (correct_count / (correct_count + wrong_count)) * 100\n",
        "  print(\"%.1f%% correct (N=%d, %d unknown)\" % (correct_percentage, (correct_count + wrong_count), discarded_count))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V3hHY0-eV6yj"
      },
      "source": [
        "If you'd like to manually evaluate particular images you can uncomment, update, and run the below cell and select an image from the test folder."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ocE3kudZq24U"
      },
      "source": [
        "# TEST_IMAGE = # UPDATE ME e.g., \"test/0/1.png\"\n",
        "# index, score = predict_image(model, TEST_IMAGE)\n",
        "# print(index, score) # prints the guessed index and the confidence"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wy80qj0XVYyi"
      },
      "source": [
        "## Generate a TensorFlow Lite Model\r\n",
        "\r\n",
        "Convert the frozen graph into a TensorFlow Lite model, which is fully quantized for use with embedded devices. The following cell will also print the model size."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "t-hU8aU24gbL"
      },
      "source": [
        "converter = tf.lite.TFLiteConverter.from_saved_model(SAVED_MODEL_FILENAME)\n",
        "model_no_quant_tflite = converter.convert()\n",
        "\n",
        "# Save the model to disk\n",
        "open(FLOAT_TFL_MODEL_FILENAME, \"wb\").write(model_no_quant_tflite)\n",
        "\n",
        "def representative_dataset():\n",
        "  for filename in glob.glob(TEST_DIR + \"/*/*.png\"):\n",
        "    img = keras.preprocessing.image.load_img(filename, target_size=(IMAGE_WIDTH, IMAGE_HEIGHT))\n",
        "    img_array = keras.preprocessing.image.img_to_array(img)\n",
        "    img_array = tf.expand_dims(img_array, 0)  # Create batch axis for images, labels in train_ds.take(1):\n",
        "    yield([img_array])\n",
        "# Set the optimization flag.\n",
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "# Enforce integer only quantization\n",
        "converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]\n",
        "converter.inference_input_type = tf.int8\n",
        "converter.inference_output_type = tf.int8\n",
        "# Provide a representative dataset to ensure we quantize correctly.\n",
        "converter.representative_dataset = representative_dataset\n",
        "model_tflite = converter.convert()\n",
        "\n",
        "# Save the model to disk\n",
        "open(QUANTIZED_TFL_MODEL_FILENAME, \"wb\").write(model_tflite)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zxV0oM8FVoww"
      },
      "source": [
        "Compare the sizes of the Tensorflow, TensorFlow Lite and Quantized TensorFlow Lite models."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NTjGMU8BPpoz"
      },
      "source": [
        "def get_dir_size(dir):\n",
        "  size = 0\n",
        "  for f in os.scandir(dir):\n",
        "    if f.is_file():\n",
        "      size += f.stat().st_size\n",
        "    elif f.is_dir():\n",
        "      size += get_dir_size(f.path)\n",
        "  return size\n",
        "\n",
        "# Calculate size\n",
        "size_tf = get_dir_size(SAVED_MODEL_FILENAME)\n",
        "size_no_quant_tflite = os.path.getsize(FLOAT_TFL_MODEL_FILENAME)\n",
        "size_tflite = os.path.getsize(QUANTIZED_TFL_MODEL_FILENAME)\n",
        "\n",
        "# Compare size\n",
        "pd.DataFrame.from_records(\n",
        "    [[\"TensorFlow\", f\"{size_tf} bytes\", \"\"],\n",
        "     [\"TensorFlow Lite\", f\"{size_no_quant_tflite} bytes \", f\"(reduced by {size_tf - size_no_quant_tflite} bytes)\"],\n",
        "     [\"TensorFlow Lite Quantized\", f\"{size_tflite} bytes\", f\"(reduced by {size_no_quant_tflite - size_tflite} bytes)\"]],\n",
        "     columns = [\"Model\", \"Size\", \"\"], index=\"Model\")\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xHDhEglPWlwE"
      },
      "source": [
        "## Test your TensorFlow Lite Models\r\n",
        "\r\n",
        "Lets now test out the TFLite models (quantized and unquantized) on the test dataset. We'll print out any gesture we get wrong as well as the percentage of known gestures correct as well as the number of gestures that were marked as unknown."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "w5QZTfwRLFAi"
      },
      "source": [
        "def predict_tflite(tflite_model, filename):\n",
        "  img = keras.preprocessing.image.load_img(filename, target_size=(IMAGE_WIDTH, IMAGE_HEIGHT))\n",
        "  img_array = keras.preprocessing.image.img_to_array(img)\n",
        "  img_array = tf.expand_dims(img_array, 0)\n",
        "\n",
        "  # Initialize the TFLite interpreter\n",
        "  interpreter = tf.lite.Interpreter(model_content=tflite_model)\n",
        "  interpreter.allocate_tensors()\n",
        "\n",
        "  input_details = interpreter.get_input_details()[0]\n",
        "  output_details = interpreter.get_output_details()[0]\n",
        "\n",
        "  # If required, quantize the input layer (from float to integer)\n",
        "  input_scale, input_zero_point = input_details[\"quantization\"]\n",
        "  if (input_scale, input_zero_point) != (0.0, 0):\n",
        "    img_array = np.multiply(img_array, 1.0 / input_scale) + input_zero_point\n",
        "    img_array = img_array.astype(input_details[\"dtype\"])\n",
        "  \n",
        "  # Invoke the interpreter\n",
        "  interpreter.set_tensor(input_details[\"index\"], img_array)\n",
        "  interpreter.invoke()\n",
        "  pred = interpreter.get_tensor(output_details[\"index\"])[0]\n",
        "  \n",
        "  # If required, dequantized the output layer (from integer to float)\n",
        "  output_scale, output_zero_point = output_details[\"quantization\"]\n",
        "  if (output_scale, output_zero_point) != (0.0, 0):\n",
        "    pred = pred.astype(np.float32)\n",
        "    pred = np.multiply((pred - output_zero_point), output_scale)\n",
        "  \n",
        "  predicted_label_index = np.argmax(pred)\n",
        "  predicted_score = pred[predicted_label_index]\n",
        "  return (predicted_label_index, predicted_score)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jdNgTO19PRqO"
      },
      "source": [
        "def run_tflite_test(model_file):\n",
        "  correct_count = 0\n",
        "  wrong_count = 0\n",
        "  discarded_count = 0\n",
        "  for label_dir in glob.glob(TEST_DIR + \"/*\"):\n",
        "    label = label_dir.replace(TEST_DIR + \"/\", \"\")\n",
        "    print(\"Testing Gesture: \",label,\" with datasize: \",len(glob.glob(label_dir + \"/*.png\")))\n",
        "    for filename in glob.glob(label_dir + \"/*.png\"):\n",
        "      index, score = predict_tflite(model_file, filename)\n",
        "      if score < 0.75:\n",
        "        discarded_count += 1\n",
        "        continue\n",
        "      if index == labelToInt[label]:\n",
        "        correct_count += 1\n",
        "      else:\n",
        "        wrong_count += 1\n",
        "        print(\"[%s] expected, [%s] found with score [%f]\" % (label, intToLabel[index], score))\n",
        "        display(Image(filename=filename))\n",
        "\n",
        "  correct_percentage = (correct_count / (correct_count + wrong_count)) * 100\n",
        "\n",
        "  print(\"%.1f%% correct (N=%d, %d unknown)\" % (correct_percentage, (correct_count + wrong_count), discarded_count))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zsJdL4dqX97n"
      },
      "source": [
        "First test the float model."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "catNwzzHX5W9"
      },
      "source": [
        "run_tflite_test(model_no_quant_tflite)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C0JDZ-5VYBNb"
      },
      "source": [
        "Then test the quantized model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uhbvbFXXYDfz"
      },
      "source": [
        "run_tflite_test(model_tflite)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7GZZJ9iKYPaX"
      },
      "source": [
        "If you'd like to manually evaluate particular images you can uncomment, update, and run the below cell and select an image from the test folder."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vtee_WxPMgup"
      },
      "source": [
        "# TEST_IMAGE = # UPDATE ME e.g., \"test/0/1.png\"\r\n",
        "# index, score = predict_tflite(model_no_quant_tflite, TEST_IMAGE)\r\n",
        "# print(\"Float model result:\")\r\n",
        "# print(index, score) # prints the guessed index and the confidence\r\n",
        "# index, score = predict_tflite(model_tflite, TEST_IMAGE)\r\n",
        "# print(\"Quantized model result:\")\r\n",
        "# print(index, score) # prints the guessed index and the confidence"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d8R-nlBYYjrP"
      },
      "source": [
        "## Generate a TensorFlow Lite for Microcontrollers Model\r\n",
        "To convert the TensorFlow Lite quantized model into a C source file that can be loaded by TensorFlow Lite for Microcontrollers on Arduino we simply need to use the ```xxd``` tool to convert the ```.tflite``` file into a ```.cc``` file."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mrvnEJLfR8KU"
      },
      "source": [
        "# Convert to a C source file, i.e, a TensorFlow Lite for Microcontrollers model\n",
        "!xxd -i {QUANTIZED_TFL_MODEL_FILENAME} > {TFL_CC_MODEL_FILENAME}\n",
        "# Update variable names\n",
        "REPLACE_TEXT = QUANTIZED_TFL_MODEL_FILENAME.replace('/', '_').replace('.', '_')\n",
        "!sed -i 's/'{REPLACE_TEXT}'/g_magic_wand_model_data/g' {TFL_CC_MODEL_FILENAME}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XfbwHwBjZL3U"
      },
      "source": [
        "That's it! You've successfully converted your TensorFlow Lite model into a TensorFlow Lite for Microcontrollers model! Run the cell below to print out its contents which we'll need for our next step, deploying the model using the Arudino IDE!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oazLUtBqWzdJ"
      },
      "source": [
        "# Print the C source file\n",
        "!cat {TFL_CC_MODEL_FILENAME}\n",
        "# !tail {TFL_CC_MODEL_FILENAME} # run this command to just see the end of the file (aka the size)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "78NF5Q1iZNsy"
      },
      "source": [
        "To download your model for use at a later date:\r\n",
        "\r\n",
        "1. On the left of the UI click on the folder icon\r\n",
        "2. Click on the three dots to the right of the ```.cc``` file you just generated and select \"download.\" The file can be found at ```models/{TFL_CC_MODEL_FILENAME}``` which by default is ```models/magic_wand.cc```\r\n",
        "\r\n",
        "Next we'll deploy that model using the Arduino IDE."
      ]
    }
  ]
}