convert_pretrained_handtracking_graph.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "convert_pretrained_handtracking_graph.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyO+GmYmRc+V86DMzPk0gcBM",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ClearSpear/cameramouse/blob/master/convert_pretrained_handtracking_graph.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ARKhzP4yK6WP",
        "colab_type": "text"
      },
      "source": [
        "## Mount Google Drive to save files"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WEJtR-odLB3e",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K5-YtTMCLDdv",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%cd '/content/gdrive/My Drive/Stanford/EE292D/project/colab_pretraineddetection'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M2CKvwvpMIoe",
        "colab_type": "text"
      },
      "source": [
        "## Prepare TF"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OWh9HfyHNoeJ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%tensorflow_version 1.x"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xGBUM723Nv47",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import tensorflow as tf\n",
        "device_name = tf.test.gpu_device_name()\n",
        "if device_name != '/device:GPU:0':\n",
        "   raise SystemError('GPU device not found')\n",
        "print('Found GPU at: {}'.format(device_name))\n",
        "print(tf.__version__)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "g9rmRrGbN1fy",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "tf.enable_eager_execution()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X6p8ozE8NwnR",
        "colab_type": "text"
      },
      "source": [
        "## Prepare TF object detection API"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qUE2t1XVN3No",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!git clone https://github.com/tensorflow/models.git"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hA5PrQZNOtR2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install tf_slim"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "g-F2VQmON_GF",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%cd /content/gdrive/My Drive/Stanford/EE292D/project/colab_pretraineddetection/models/research/\n",
        "!protoc object_detection/protos/*.proto --python_out=."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cNThUamCOqkj",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "os.environ['PYTHONPATH'] += ':/content/gdrive/My Drive/Stanford/EE292D/project/colab/models/research/:/content/gdrive/My Drive/Stanford/EE292D/project/colab/models/research/slim'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c4-n8QIAOsU7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!python setup.py build\n",
        "!python setup.py install"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Aj8D6g4lO4cc",
        "colab_type": "text"
      },
      "source": [
        "Test that installation was successful"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NP7mkA3cO6Nt",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%cd /content/gdrive/My Drive/Stanford/EE292D/project/colab_pretraineddetection/models/research/object_detection/builders/\n",
        "!python model_builder_test.py"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d3HBRPcQKxv6",
        "colab_type": "text"
      },
      "source": [
        "## Download the github with the trained handtracking graph"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qU4XPyS1Ktnw",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!git clone https://github.com/victordibia/handtracking.git"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KWAIs2XRLjo2",
        "colab_type": "text"
      },
      "source": [
        "## Convert saved model to frozen graph compatible with tflite"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E52jvMYkLlys",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%cd /content/gdrive/My Drive/Stanford/EE292D/project/colab_pretraineddetection\n",
        "H = 300\n",
        "W = 300"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F4R0lrExLp_Z",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!python models/research/object_detection/export_tflite_ssd_graph.py \\\n",
        "--pipeline_config_path=handtracking/model-checkpoint/ssdmobilenetv1/ssd_mobilenet_v1_coco.config \\\n",
        "--trained_checkpoint_prefix=handtracking/model-checkpoint/ssdmobilenetv1/model.ckpt-200002 \\\n",
        "--output_directory=tflite \\\n",
        "--add_postprocessing_op=true"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "De-EHQI_LyRQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "gf = tf.GraphDef()   \n",
        "m_file = open('tflite/tflite_graph.pb','rb')\n",
        "gf.ParseFromString(m_file.read())\n",
        "\n",
        "for n in gf.node:\n",
        "    print( n.name )\n",
        "    tensor = n.op # save the input and output tensor"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oNpXldvJL2s5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# A generator that provides a representative dataset\n",
        "def representative_data_gen():\n",
        "  dataset_list = tf.data.Dataset.list_files('/content/gdrive/My Drive/Stanford/EE292D/project/colab/egohands/*/*')\n",
        "  for i in range(100):\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, [H, W])\n",
        "    image = tf.cast(image / 255., tf.float32)\n",
        "    image = tf.expand_dims(image, 0)\n",
        "    yield [image]\n",
        "\n",
        "input_tensors = ['normalized_input_image_tensor']\n",
        "output_tensors = ['TFLite_Detection_PostProcess','TFLite_Detection_PostProcess:1','TFLite_Detection_PostProcess:2','TFLite_Detection_PostProcess:3']\n",
        "\n",
        "converter = tf.lite.TFLiteConverter.from_frozen_graph('tflite/tflite_graph.pb', input_tensors, output_tensors, input_shapes={'normalized_input_image_tensor':[1,H,W,3]})\n",
        "converter.allow_custom_ops = True\n",
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\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",
        "converter.representative_dataset = representative_data_gen\n",
        "tflite_model = converter.convert()\n",
        "\n",
        "with open(\"handdetection_ssdmobilenetv1.tflite\", \"wb\") as f:\n",
        "  f.write(tflite_model)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}