colab-notebooks/Easy_Object_Detection_With_Custom_Data_Demo_Training.ipynb

{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "mJz_ToJtkufM"
      },
      "source": [
        "# DepthAI Tutorial: Training an Object Detection Model with Your Own Data\n",
        "\n",
        "<img src=\"https://docs.luxonis.com/images/depthai_logo.png\" width=\"500\">\n",
        "\n",
        "**<font color=\"red\">Warning! This tutorial is now deprecated. Please refer to [YoloV5 tutorial](https://colab.research.google.com/github/luxonis/depthai-ml-training/blob/master/colab-notebooks/YoloV5_training.ipynb), [YoloV6 tutorial](https://colab.research.google.com/github/luxonis/depthai-ml-training/blob/master/colab-notebooks/YoloV6_training.ipynb) or [YoloV7 tutorial](https://colab.research.google.com/github/luxonis/depthai-ml-training/blob/master/colab-notebooks/YoloV7_training.ipynb).</font>**\n",
        "\n",
        "Welcome to DepthAI! \n",
        "\n",
        "In this tutorial we will go through the basic training of an object detection model with your own annotated images.\n",
        "\n",
        "The model you will use is a pretrained Mobilenet SSD v2 from the Tensorflow Object Detection API model zoo. The framework used for training is TensorFlow 1.15.2.\n",
        "Will run through the following steps:\n",
        "\n",
        "\n",
        "*   Install the libraries\n",
        "*   Clone the github repo and replace the repo training data with your data from google drive\n",
        "    (if you do not have own data, please see this\n",
        "*   Train the model on the new images\n",
        "*   Run inference on a few images to see what the model can detect\n",
        "*   Convert the model to OpenVINO Intermediate Representation\n",
        "*   To run the model on DepthAI modules, compile the IR obtained above to a .blob file \n",
        "\n",
        "You can make a copy of this tutorial: File-> Save a copy in Drive\n",
        "\n",
        "Note: the model training can be run with the repo images of three fruits if you choose to skip the part with loading your own images\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Kmw4ac-EQuzj"
      },
      "source": [
        "# Install Libraries\n",
        "# Press Shift+Enter to run cells\n",
        "Some cells are commented out (all text is green), so they will not run the code inside.\n",
        "To uncomment code inside a cell, select all the code then press 'Ctrl' + '/'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JVmWjRMLlR3l"
      },
      "source": [
        "###After this cell executes runtime will restart to complete the install, ignore and close the message, continue running the cells below this one\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pH1x08R-yM-L"
      },
      "outputs": [],
      "source": [
        "# %%capture\n",
        "#After this cell executes runtime will restart to finish the install, ignore and close the crash message, continue running cells starting with the one below\n",
        "!pip install numpy==1.17.5;"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9z8kgcs0E7iV"
      },
      "outputs": [],
      "source": [
        "!pip install tf_slim==1.1.0\n",
        "!pip install tensorflow-gpu==1.15"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gnNXNQCjdniL"
      },
      "outputs": [],
      "source": [
        "# For faster training time, images should be resized to 300x300 and then annotated\n",
        "# Images should contain the objects of interest at various scales, angles, lighting conditions, locations\n",
        "# For acceptable results - mAP@0.5 of 0.9 the model was trained with batch size of 24\n",
        "# and 5000 steps. this takes about 1h using 2 augmentations. \n",
        "# using 5 augmentations it takes about 2h \n",
        "num_steps = 5000  # A step means using a single batch of data. larger batch, less steps required\n",
        "#Number of evaluation steps.\n",
        "num_eval_steps = 50\n",
        "#Batch size 24 is a setting that generally works well. can be changed higher or lower \n",
        "MODELS_CONFIG = {\n",
        "        'ssd_mobilenet_v2': {\n",
        "        'model_name': 'ssd_mobilenet_v2_coco_2018_03_29',\n",
        "        'pipeline_file': 'ssd_mobilenet_v2_coco.config',\n",
        "        'batch_size': 24\n",
        "    }\n",
        "}\n",
        "selected_model = 'ssd_mobilenet_v2'\n",
        "\n",
        "# Name of the object detection model to use.\n",
        "MODEL = MODELS_CONFIG[selected_model]['model_name']\n",
        "\n",
        "# Name of the pipline file in tensorflow object detection API.\n",
        "pipeline_file = MODELS_CONFIG[selected_model]['pipeline_file']\n",
        "\n",
        "# Training batch size fits in Colab's GPU memory for selected model.\n",
        "batch_size = MODELS_CONFIG[selected_model]['batch_size']"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w4V-XE6kbkc1"
      },
      "source": [
        "## Clone the `depthai-ml-training` repository"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dxc3DmvLQF3z"
      },
      "outputs": [],
      "source": [
        "repo_url = 'https://github.com/luxonis/depthai-ml-training.git'\n",
        "import os\n",
        "%cd /content\n",
        "repo_dir_path = \"/content/depthai-ml-training\"\n",
        "!rm -rf {repo_dir_path}\n",
        "!git clone {repo_url} {repo_dir_path}\n",
        "%cd {repo_dir_path}\n",
        "!git pull"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BSJEHMWmltIl"
      },
      "source": [
        "## Download our example dataset of PASCAL cats, dogs, and birds\n",
        "\n",
        "You can skip this if you are using your own data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "13dDLE4-l6oa"
      },
      "outputs": [],
      "source": [
        "%cd /content\n",
        "!pip install -U --no-cache-dir gdown==4.4.0 --pre\n",
        "!gdown --id 1Sr3E0O10sMP8YmhCDHWTqD2Wb7EbMuG-\n",
        "!unzip \"/content/pascal_animals.zip\"\n",
        "!rm -rf \"/content/pascal_animals.zip\"\n",
        "!mv /content/sorted/* /content\n",
        "!rm -rf sorted/"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6LnZjgF4N71_"
      },
      "source": [
        "#[Optional] If you wish to use your own data, run the cells below.\n",
        "(Otherwise the model can train with the data already present in the cloned repo)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hPVANnM4QUWT"
      },
      "source": [
        "## Mount your google drive"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KPahiXuPNQwX"
      },
      "outputs": [],
      "source": [
        "#mount your google drive.\n",
        "#it will be visible in the file navigator on the left of this notebook\n",
        "#there should be a folder in your drive with your data\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a2ZtJqdRsV0A"
      },
      "source": [
        "Own data should be split into 80% training (train folder), 20% for validation (test folder). The train and test folder should contain images and associated .xml annotations. \n",
        "\n",
        "Have extra images for final model testing (final_test folder). These need not be annotated.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gsUxoYiDOYwr"
      },
      "outputs": [],
      "source": [
        "#copy files from gdrive to colab drive. this takes a few minutes, depending on the number of files.\n",
        "#Go on the file explorer on the left of this notebook and access your gdrive. find the folders\n",
        "#with your train, test and final_test images.\n",
        "#Right click on each and copy the path. paste it btw the first \" \" in the corresponding lines  \n",
        "\n",
        "#training folder\n",
        "# !cp -r \"/path/to/data/on/Google/Drive/train\" \"/content/train\"\n",
        "#validation folder\n",
        "# !cp -r \"/path/to/data/on/Google/Drive/validation\" \"/content/validation\"\n",
        "#test folder\n",
        "# !cp -r \"/path/to/data/on/Google/Drive/test\" \"/content/test\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pHUcVDSEfuCE"
      },
      "outputs": [],
      "source": [
        "# quick check for training data files. you can also browse to the object_detection_demo_flows\n",
        "# on the left and see if they were copied\n",
        "!ls /content/train"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bI8__uNS8-ns"
      },
      "source": [
        "# Install Tensorflow Object Detection API\n",
        "\n",
        "Clone TF models which contains the Object Detection API; also install the required dependencies\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ecpHEnka8Kix"
      },
      "outputs": [],
      "source": [
        "# %%capture\n",
        "%cd /content\n",
        "!git clone --quiet https://github.com/tensorflow/models.git\n",
        "%cd /content/models/\n",
        "!git checkout 58d19c67e1d30d905dd5c6e5092348658fed80af\n",
        "!apt-get update && apt-get install -y -qq protobuf-compiler python-pil python-lxml python-tk\n",
        "!pip install -q Cython==0.29.28 contextlib2==0.5.5 pillow==7.1.2 lxml==4.2.6 matplotlib==3.2.2\n",
        "!pip uninstall -y pycocotools\n",
        "!pip install --no-cache-dir pycocotools==2.0.4\n",
        "%cd /content/models/research\n",
        "!protoc object_detection/protos/*.proto --python_out=.\n",
        "import os\n",
        "os.environ['PYTHONPATH'] += ':/content/models/research/:/content/models/research/slim/'\n",
        "!python object_detection/builders/model_builder_test.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u-k7uGThXlny"
      },
      "source": [
        "## Prepare `tfrecord` files\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ezGDABRXXhPP"
      },
      "outputs": [],
      "source": [
        "# %%capture\n",
        "%cd /content/\n",
        "\n",
        "# Convert train folder annotation xml files to a single csv file,\n",
        "# generate the `label_map.pbtxt` file to `data/` directory as well.\n",
        "!python depthai-ml-training/helpers/xml_to_csv.py -i /content/train/ -o /content/train_labels.csv -l /content\n",
        "# Convert test folder annotation xml files to a single csv.\n",
        "!python depthai-ml-training/helpers/xml_to_csv.py -i /content/validation/ -o /content/test_labels.csv\n",
        "\n",
        "# Generate `train.record`\n",
        "!python depthai-ml-training/helpers/generate_tfrecord.py --csv_input=/content/train_labels.csv --output_path=/content/train.record --img_path=/content/train --label_map label_map.pbtxt\n",
        "\n",
        "# Generate `test.record`\n",
        "!python depthai-ml-training/helpers/generate_tfrecord.py --csv_input=/content/test_labels.csv --output_path=/content/test.record --img_path=/content/validation --label_map label_map.pbtxt\n",
        "\n",
        "# Set the paths\n",
        "test_record_fname = '/content/test.record'\n",
        "train_record_fname = '/content/train.record'\n",
        "label_map_pbtxt_fname = '/content/label_map.pbtxt'\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iCNYAaC7w6N8"
      },
      "source": [
        "## Download the Mobilenet SSD v2 Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "orDCj6ihgUMR"
      },
      "outputs": [],
      "source": [
        "%cd /content/models/research\n",
        "\n",
        "import os\n",
        "import shutil\n",
        "import glob\n",
        "import urllib.request\n",
        "import tarfile\n",
        "MODEL_FILE = MODEL + '.tar.gz'\n",
        "DOWNLOAD_BASE = 'http://download.tensorflow.org/models/object_detection/'\n",
        "DEST_DIR = '/content/models/research/pretrained_model'\n",
        "\n",
        "if not (os.path.exists(MODEL_FILE)):\n",
        "    urllib.request.urlretrieve(DOWNLOAD_BASE + MODEL_FILE, MODEL_FILE)\n",
        "\n",
        "tar = tarfile.open(MODEL_FILE)\n",
        "tar.extractall()\n",
        "tar.close()\n",
        "\n",
        "os.remove(MODEL_FILE)\n",
        "if (os.path.exists(DEST_DIR)):\n",
        "    shutil.rmtree(DEST_DIR)\n",
        "os.rename(MODEL, DEST_DIR)\n",
        "!echo {DEST_DIR}\n",
        "!ls -alh {DEST_DIR}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UHnxlfRznPP3"
      },
      "outputs": [],
      "source": [
        "#TF pretrained model checkpoint\n",
        "fine_tune_checkpoint = os.path.join(DEST_DIR, \"model.ckpt\")\n",
        "fine_tune_checkpoint"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MvwtHlLOeRJD"
      },
      "source": [
        "## Configuring a Training Pipeline"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dIhw7IdpLuiU"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "pipeline_fname = os.path.join('/content/models/research/object_detection/samples/configs/', pipeline_file)\n",
        "\n",
        "assert os.path.isfile(pipeline_fname), '`{}` not exist'.format(pipeline_fname)\n",
        "def get_num_classes(pbtxt_fname):\n",
        "    from object_detection.utils import label_map_util\n",
        "    label_map = label_map_util.load_labelmap(pbtxt_fname)\n",
        "    categories = label_map_util.convert_label_map_to_categories(\n",
        "        label_map, max_num_classes=90, use_display_name=True)\n",
        "    category_index = label_map_util.create_category_index(categories)\n",
        "    return len(category_index.keys())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FRofNAq5GDXY"
      },
      "outputs": [],
      "source": [
        "from pathlib import Path\n",
        "parent = Path(label_map_pbtxt_fname).parent\n",
        "!ls {parent}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YjtCbLF2i0wI"
      },
      "outputs": [],
      "source": [
        "import re\n",
        "iou_threshold = 0.50\n",
        "num_classes = get_num_classes(label_map_pbtxt_fname)\n",
        "with open(pipeline_fname) as f:\n",
        "    s = f.read()\n",
        "with open(pipeline_fname, 'w') as f:\n",
        "    \n",
        "    # fine_tune_checkpoint\n",
        "    s = re.sub('fine_tune_checkpoint: \".*?\"',\n",
        "               'fine_tune_checkpoint: \"{}\"'.format(fine_tune_checkpoint), s)\n",
        "    \n",
        "    # tfrecord files train and test.\n",
        "    s = re.sub(\n",
        "        '(input_path: \".*?)(train.record)(.*?\")', 'input_path: \"{}\"'.format(train_record_fname), s)\n",
        "    s = re.sub(\n",
        "        '(input_path: \".*?)(val.record)(.*?\")', 'input_path: \"{}\"'.format(test_record_fname), s)\n",
        "\n",
        "    # label_map_path\n",
        "    s = re.sub(\n",
        "        'label_map_path: \".*?\"', 'label_map_path: \"{}\"'.format(label_map_pbtxt_fname), s)\n",
        "\n",
        "    # Set training batch_size.\n",
        "    s = re.sub('batch_size: [0-9]+',\n",
        "               'batch_size: {}'.format(batch_size), s)\n",
        "\n",
        "    # Set training steps, num_steps\n",
        "    s = re.sub('num_steps: [0-9]+',\n",
        "               'num_steps: {}'.format(num_steps), s)\n",
        "    \n",
        "    # Set number of classes num_classes.\n",
        "    s = re.sub('num_classes: [0-9]+',\n",
        "               'num_classes: {}'.format(num_classes), s)\n",
        "    # Set number of classes num_classes.\n",
        "    s = re.sub('iou_threshold: [0-9].[0-9]+',\n",
        "               'iou_threshold: {}'.format(iou_threshold), s)\n",
        "    \n",
        "    f.write(s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GH0MEEanocn6"
      },
      "outputs": [],
      "source": [
        "# #Have a look at the config file with various settings\n",
        "# !cat {pipeline_fname}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JDddx2rPfex9"
      },
      "source": [
        "# Train the model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PS6W4ZEEWBgM"
      },
      "source": [
        "##[Optional] The cell below adds Tensorboard visualization to the training process.\n",
        "Will open in new tab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f11w0uO3jFCB"
      },
      "outputs": [],
      "source": [
        "#After running this cell click on the link ending in 'ngrok.io' in the output cell to open tensorboard\n",
        "#Tensoarboard will show you graphically different training parameters as the model is training\n",
        "#when training finishes after the set number of steps, tensorboard can be used to see a nice summary of the training process\n",
        "#Visuals will load in Tensorboard after the model has gone through a few hundred steps\n",
        "\n",
        "#You may need to re-run this cell; replace any files if prompted if you get an error\n",
        "\n",
        "model_dir = 'training/'\n",
        "!wget https://bin.equinox.io/c/4VmDzA7iaHb/ngrok-stable-linux-amd64.zip\n",
        "!unzip ngrok-stable-linux-amd64.zip\n",
        "LOG_DIR = model_dir\n",
        "get_ipython().system_raw(\n",
        "    'tensorboard --logdir {} --host 0.0.0.0 --port 6006 &'\n",
        "    .format(LOG_DIR)\n",
        ")\n",
        "get_ipython().system_raw('./ngrok http 6006 &')\n",
        "! curl -s http://localhost:4040/api/tunnels | python3 -c \\\n",
        "    \"import sys, json; print(json.load(sys.stdin)['tunnels'][0]['public_url'])\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zOyWAHNywiL4"
      },
      "source": [
        "## Start the training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CjDHjhKQofT5"
      },
      "outputs": [],
      "source": [
        "model_dir = 'training/'\n",
        "# Optionally remove content in output model directory for a fresh start.\n",
        "# !rm -rf {model_dir}\n",
        "# os.makedirs(model_dir, exist_ok=True)\n",
        "!python /content/models/research/object_detection/model_main.py \\\n",
        "    --pipeline_config_path={pipeline_fname} \\\n",
        "    --model_dir={model_dir} \\\n",
        "    --alsologtostderr \\\n",
        "    --num_train_steps={num_steps} \\\n",
        "    --num_eval_steps={num_eval_steps}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KP-tUdtnRybs"
      },
      "outputs": [],
      "source": [
        "#model dir check for the trained model\n",
        "!ls {model_dir}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OmSESMetj1sa"
      },
      "source": [
        "## Export a Trained Inference Graph\n",
        "Once your training job is complete, you need to extract the newly trained inference graph, which will be later used to perform the object detection. This can be done as follows:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zZJ_AF5MD3hZ"
      },
      "outputs": [],
      "source": [
        "#clean output_directory if necessary to start fresh:\n",
        "\n",
        "# !rm -rf /content/object_detection_demo/fine_tuned_model/ \n",
        "# os.makedirs('/content/object_detection_demo_flow/fine_tuned_model/', exist_ok=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DHoP90pUyKSq"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "import re\n",
        "import numpy as np\n",
        "\n",
        "output_directory = './fine_tuned_model'\n",
        "# output_directory = '/content/gdrive/My\\ Drive/data/'\n",
        "\n",
        "lst = os.listdir(model_dir)\n",
        "lst = [l for l in lst if 'model.ckpt-' in l and '.meta' in l]\n",
        "steps=np.array([int(re.findall('\\d+', l)[0]) for l in lst])\n",
        "last_model = lst[steps.argmax()].replace('.meta', '')\n",
        "\n",
        "last_model_path = os.path.join(model_dir, last_model)\n",
        "print(last_model_path)\n",
        "!python /content/models/research/object_detection/export_inference_graph.py \\\n",
        "    --input_type=image_tensor \\\n",
        "    --pipeline_config_path={pipeline_fname} \\\n",
        "    --output_directory={output_directory} \\\n",
        "    --trained_checkpoint_prefix={last_model_path}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "usgBZvkz0nqD"
      },
      "outputs": [],
      "source": [
        "#export directory check\n",
        "!ls {output_directory}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CnDo1lonKgFr"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "pb_fname = os.path.join(os.path.abspath(output_directory), \"frozen_inference_graph.pb\")\n",
        "assert os.path.isfile(pb_fname), '`{}` not exist'.format(pb_fname)\n",
        "# !ls -alh {pb_fname}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mz1gX19GlVW7"
      },
      "source": [
        "## Running Inference: Checking what the trained model can detect\n",
        "Test with images in repository `object_detection_demo_flow/data/images/final test` directory."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Pzj9A4e5mj5l"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import glob\n",
        "\n",
        "# Path to frozen detection graph. This is the actual model that is used for the object detection.\n",
        "PATH_TO_CKPT = pb_fname\n",
        "\n",
        "# List of the strings that is used to add correct label for each box.\n",
        "PATH_TO_LABELS = label_map_pbtxt_fname\n",
        "\n",
        "# If you want to test the code with your images, just add images files to the PATH_TO_TEST_IMAGES_DIR.\n",
        "PATH_TO_TEST_IMAGES_DIR =  \"/content/test\"\n",
        "\n",
        "assert os.path.isfile(pb_fname)\n",
        "assert os.path.isfile(PATH_TO_LABELS)\n",
        "TEST_IMAGE_PATHS = glob.glob(os.path.join(PATH_TO_TEST_IMAGES_DIR, \"*.*\"))\n",
        "assert len(TEST_IMAGE_PATHS) > 0, 'No image found in `{}`.'.format(PATH_TO_TEST_IMAGES_DIR)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CG5YUMdg1Po7"
      },
      "outputs": [],
      "source": [
        "%cd /content/models/research/object_detection\n",
        "\n",
        "import numpy as np\n",
        "import os\n",
        "import six.moves.urllib as urllib\n",
        "import sys\n",
        "import tarfile\n",
        "import tensorflow as tf\n",
        "import zipfile\n",
        "\n",
        "from collections import defaultdict\n",
        "from io import StringIO\n",
        "# This is needed to display the images.\n",
        "%matplotlib inline\n",
        "import matplotlib.pyplot as plt\n",
        "from PIL import Image\n",
        "\n",
        "# This is needed since the notebook is stored in the object_detection folder.\n",
        "sys.path.append(\"..\")\n",
        "from object_detection.utils import ops as utils_ops\n",
        "\n",
        "from object_detection.utils import label_map_util\n",
        "\n",
        "from object_detection.utils import visualization_utils as vis_util\n",
        "\n",
        "\n",
        "detection_graph = tf.Graph()\n",
        "with detection_graph.as_default():\n",
        "    od_graph_def = tf.GraphDef()\n",
        "    with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:\n",
        "        serialized_graph = fid.read()\n",
        "        od_graph_def.ParseFromString(serialized_graph)\n",
        "        tf.import_graph_def(od_graph_def, name='')\n",
        "\n",
        "\n",
        "label_map = label_map_util.load_labelmap(PATH_TO_LABELS)\n",
        "categories = label_map_util.convert_label_map_to_categories(\n",
        "    label_map, max_num_classes=num_classes, use_display_name=True)\n",
        "category_index = label_map_util.create_category_index(categories)\n",
        "\n",
        "\n",
        "def load_image_into_numpy_array(image):\n",
        "    (im_width, im_height) = image.size\n",
        "    return np.array(image.getdata()).reshape(\n",
        "        (im_height, im_width, 3)).astype(np.uint8)\n",
        "\n",
        "# Size, in inches, of the output images.\n",
        "IMAGE_SIZE = (12, 8)\n",
        "\n",
        "\n",
        "def run_inference_for_single_image(image, graph):\n",
        "    with graph.as_default():\n",
        "        with tf.Session() as sess:\n",
        "            # Get handles to input and output tensors\n",
        "            ops = tf.get_default_graph().get_operations()\n",
        "            all_tensor_names = {\n",
        "                output.name for op in ops for output in op.outputs}\n",
        "            tensor_dict = {}\n",
        "            for key in [\n",
        "                'num_detections', 'detection_boxes', 'detection_scores',\n",
        "                'detection_classes', 'detection_masks'\n",
        "            ]:\n",
        "                tensor_name = key + ':0'\n",
        "                if tensor_name in all_tensor_names:\n",
        "                    tensor_dict[key] = tf.get_default_graph().get_tensor_by_name(\n",
        "                        tensor_name)\n",
        "            if 'detection_masks' in tensor_dict:\n",
        "                # The following processing is only for single image\n",
        "                detection_boxes = tf.squeeze(\n",
        "                    tensor_dict['detection_boxes'], [0])\n",
        "                detection_masks = tf.squeeze(\n",
        "                    tensor_dict['detection_masks'], [0])\n",
        "                # Reframe is required to translate mask from box coordinates to image coordinates and fit the image size.\n",
        "                real_num_detection = tf.cast(\n",
        "                    tensor_dict['num_detections'][0], tf.int32)\n",
        "                detection_boxes = tf.slice(detection_boxes, [0, 0], [\n",
        "                                           real_num_detection, -1])\n",
        "                detection_masks = tf.slice(detection_masks, [0, 0, 0], [\n",
        "                                           real_num_detection, -1, -1])\n",
        "                detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks(\n",
        "                    detection_masks, detection_boxes, image.shape[0], image.shape[1])\n",
        "                detection_masks_reframed = tf.cast(\n",
        "                    tf.greater(detection_masks_reframed, 0.5), tf.uint8)\n",
        "                # Follow the convention by adding back the batch dimension\n",
        "                tensor_dict['detection_masks'] = tf.expand_dims(\n",
        "                    detection_masks_reframed, 0)\n",
        "            image_tensor = tf.get_default_graph().get_tensor_by_name('image_tensor:0')\n",
        "\n",
        "            # Run inference\n",
        "            output_dict = sess.run(tensor_dict,\n",
        "                                   feed_dict={image_tensor: np.expand_dims(image, 0)})\n",
        "\n",
        "            # all outputs are float32 numpy arrays, so convert types as appropriate\n",
        "            output_dict['num_detections'] = int(\n",
        "                output_dict['num_detections'][0])\n",
        "            output_dict['detection_classes'] = output_dict[\n",
        "                'detection_classes'][0].astype(np.uint8)\n",
        "            output_dict['detection_boxes'] = output_dict['detection_boxes'][0]\n",
        "            output_dict['detection_scores'] = output_dict['detection_scores'][0]\n",
        "            if 'detection_masks' in output_dict:\n",
        "                output_dict['detection_masks'] = output_dict['detection_masks'][0]\n",
        "    return output_dict"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QGFsIDVwqYm-"
      },
      "outputs": [],
      "source": [
        "# Visualize the results\n",
        "%matplotlib inline\n",
        "plt.close('all')\n",
        "for image_path in TEST_IMAGE_PATHS[:21]:\n",
        "  image = Image.open(image_path)\n",
        "  print(image_path)\n",
        "  # the array based representation of the image will be used later in order to prepare the\n",
        "  # result image with boxes and labels on it.\n",
        "  image_np = load_image_into_numpy_array(image)\n",
        "  # Expand dimensions since the model expects images to have shape: [1, None, None, 3]\n",
        "  image_np_expanded = np.expand_dims(image_np, axis=0)\n",
        "  # Actual detection.\n",
        "  output_dict = run_inference_for_single_image(image_np, detection_graph)\n",
        "  # Visualization of the results of a detection.\n",
        "  vis_util.visualize_boxes_and_labels_on_image_array(\n",
        "      image_np,\n",
        "      output_dict['detection_boxes'],\n",
        "      output_dict['detection_classes'],\n",
        "      output_dict['detection_scores'],\n",
        "      category_index,\n",
        "      instance_masks=output_dict.get('detection_masks'),\n",
        "      use_normalized_coordinates=True,\n",
        "      line_thickness=8)\n",
        "  plt.figure(figsize=IMAGE_SIZE)\n",
        "  plt.imshow(image_np)\n",
        "  plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_PlfZAR1OCK2"
      },
      "source": [
        "# Convert TF model to OpenVINO 21.03 Intermediate Representation (IR)\n",
        " This can be used to run inference on OpenVINO.\n",
        "# In order to run the model on DepthAI modules, we then compile the IR obtained above to a .blob (via a server we set up just for that) "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3qucQJwsWhL3"
      },
      "source": [
        "## First, we install Open Vino 21.03\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0BxCxkQ6NkcF"
      },
      "outputs": [],
      "source": [
        "%cd ../.."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5py5tfxS6VaA"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from urllib.parse import urlparse\n",
        "\n",
        "## install tools. Open Vino takes some time to download - it's ~400MB\n",
        "!sudo apt-get install -y pciutils cpio\n",
        "!sudo apt autoremove\n",
        "\n",
        "## downnload installation files\n",
        "url = \"https://registrationcenter-download.intel.com/akdlm/irc_nas/17662/l_openvino_toolkit_p_2021.3.394.tgz\"\n",
        "!wget {url}\n",
        "\n",
        "## Get the name of the tgz\n",
        "parsed = urlparse(url)\n",
        "openvino_tgz = os.path.basename(parsed.path)\n",
        "openvino_folder = os.path.splitext(openvino_tgz)[0]\n",
        "\n",
        "## Extract & install openvino\n",
        "!tar xf {openvino_tgz}\n",
        "%cd {openvino_folder}\n",
        "!./install_openvino_dependencies.sh && \\\n",
        "    sed -i 's/decline/accept/g' silent.cfg && \\\n",
        "    ./install.sh --silent silent.cfg"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r2enKxK_qakX"
      },
      "outputs": [],
      "source": [
        "!ls"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h-jbjiw-7nAv"
      },
      "source": [
        "[Optional] Open Vino install check, generally not needed"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0NJp28Oj7Ts9"
      },
      "outputs": [],
      "source": [
        "# !source /opt/intel/openvino/bin/setupvars.sh && \\\n",
        "#     /opt/intel/openvino/deployment_tools/demo/demo_squeezenet_download_convert_run.sh"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TygLw4WIO8BD"
      },
      "source": [
        "### Here we run some modifications in the ssd2 OpenVINO extension for TF so that our Mobilenet SSDv2 model can convert successfully to the IR"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P-bGZHGcMNbZ"
      },
      "outputs": [],
      "source": [
        "%cd /opt/intel/openvino_2021/deployment_tools/model_optimizer/extensions/front/tf/\n",
        "\n",
        "#openvino fixes: edit \n",
        "# Read in the file, make sure the .json corresponds to the model!!!\n",
        "with open('ssd_v2_support.json', 'r') as file :\n",
        "  filedata = file.read()\n",
        "\n",
        "# Replace the target string\n",
        "filedata = filedata.replace('\"Postprocessor/ToFloat\"', '\"Postprocessor/Cast_1\"')\n",
        "\n",
        "# Write the file out again\n",
        "with open('ssd_v2_support.json', 'w') as file:\n",
        "  file.write(filedata)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R7sqTX79W6hh"
      },
      "source": [
        "## Convert TF model to Open Vino Intermediate Representation\n",
        "If using own model, please change to your desired name for output directory --output_dir \"choose name\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LsWggE5AIWS6"
      },
      "outputs": [],
      "source": [
        "#CONVERT TF MODEL to OPEN VINO IRv10. saved in IR_V10_fruits_mnssdv2_6k directory or\n",
        "#choose own name for --output_dir \"choose name\"\n",
        "%cd \"/content/models/research/fine_tuned_model/\"\n",
        "!source /opt/intel/openvino_2021/bin/setupvars.sh && \\\n",
        "    python /opt/intel/openvino_2021/deployment_tools/model_optimizer/mo.py \\\n",
        "    --input_model frozen_inference_graph.pb \\\n",
        "    --tensorflow_use_custom_operations_config /opt/intel/openvino_2021/deployment_tools/model_optimizer/extensions/front/tf/ssd_v2_support.json \\\n",
        "    --tensorflow_object_detection_api_pipeline_config pipeline.config \\\n",
        "    --reverse_input_channels \\\n",
        "    --output_dir ./pascal_animals \\\n",
        "    --data_type FP16"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "o9R7LTkNvZos"
      },
      "outputs": [],
      "source": [
        "# from google.colab import drive\n",
        "# drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P15Hl-Fbmkib"
      },
      "outputs": [],
      "source": [
        "#check directory containing the exported TF trained model and the IRv10 folder\n",
        "%ls pascal_animals/"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WD2BMtXoPjYN"
      },
      "source": [
        "## Now we compile the IR model to a .blob for use on DepthAI modules/platform\n",
        "\n",
        "---\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZqZqajTIP3Sv"
      },
      "source": [
        "### We save the blob in the IR directory from above, corresponding to --output_dir parameter above. \n",
        "The blob filename will be *frozen_inference_graph.blob*"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IqiPpnCwPioH"
      },
      "outputs": [],
      "source": [
        "#No changes needed here unless using custom data.\n",
        "#CHOOSE the directory where you would like to save the blob.\n",
        "# I use the same --output_dir as above for the IR conversion\n",
        "blob_dir = \"/content/models/research/fine_tuned_model/pascal_animals/\"\n",
        "\n",
        "#Copy the path of your .xml and .bin files. For that, you can look at the IR\n",
        "#conversion output cell, select and copy from:\n",
        "#[SUCCESS] XML file and bin file paths.\n",
        "#Or you can choose to compile other .xml .bin files from a different location\n",
        "#\n",
        "xmlfile = \"/content/models/research/fine_tuned_model/pascal_animals/frozen_inference_graph.xml\"\n",
        "binfile = \"/content/models/research/fine_tuned_model/pascal_animals/frozen_inference_graph.bin\"\n",
        "!python -m pip install boto3==1.17.39\n",
        "!python -m pip install blobconverter==1.2.9\n",
        "\n",
        "import blobconverter\n",
        "blob_path = blobconverter.from_openvino(\n",
        "    xml=xmlfile,\n",
        "    bin=binfile,\n",
        "    data_type=\"FP16\",\n",
        "    shaves=6,\n",
        ")\n",
        "from google.colab import files\n",
        "files.download(blob_path) "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K4hQdGwQSxP0"
      },
      "source": [
        "##Now you can download your .blob file and run it on the DepthAI module/platform\n",
        "To download locally, use the file explorer on the left to locate the file in the --output_dir folder, then right click download. Colab takes a few seconds to prepare the file, then the download prompt will appear."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "24-2kaAFvyEd"
      },
      "source": [
        "# Run the blob on DepthAI\n",
        "\n",
        "To use this blob with DepthAI:\n",
        "\n",
        "- Clone depthai demo script\n",
        "\n",
        "    ```\n",
        "    $ git clone https://github.com/luxonis/depthai.git\n",
        "    ```\n",
        "- Install requirements\n",
        "\n",
        "    ```\n",
        "    $ python3 install_requirements.py\n",
        "    ```\n",
        "- Create a new model directory and insert the blob\n",
        "\n",
        "    ```\n",
        "    $ mkdir resources/nn/custom_mobilenet\n",
        "    $ mv <downloaded_blob> resources/nn/custom_mobilenet/\n",
        "    ```\n",
        "\n",
        "- Copy config file from mobilenet-ssd model\n",
        "\n",
        "    ```\n",
        "    $ cp resources/nn/mobilenet-ssd/mobilenet-ssd.json resources/nn/custom_mobilenet/custom_mobilenet.json\n",
        "    ```\n",
        "\n",
        "- Run the demo script\n",
        "\n",
        "    ```\n",
        "    $ python3 depthai_demo.py -cnn custom_mobilenet\n",
        "    ```\n",
        "\n",
        "You can read more details on deployment [here](https://docs.luxonis.com/en/latest/pages/tutorials/first_steps/#using-custom-models)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U18F9dvOLs1I"
      },
      "outputs": [],
      "source": [
        "#Compress the folder with the Tensorflow model and OpenVINO IR_V10 folder and download it locally.\n",
        "#These are in content/models/research/fine_tuned_model, so we compress the fine_tuned_model folder.\n",
        "!tar czvf fine_tuned_model.tar.gz {blob_dir}\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YARGCLW7KJxo"
      },
      "outputs": [],
      "source": [
        "#download the compressed IRv10 folder locally\n",
        "#or can use file navigator on the left to move it to your gdrive\n",
        "from google.colab import files\n",
        "files.download(\"fine_tuned_model.tar.gz\")"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "6LnZjgF4N71_",
        "PS6W4ZEEWBgM"
      ],
      "name": "Update_Easy_Object_Detection_With_Custom_Data_Demo_Training_Git.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.10.6 (main, Nov  2 2022, 18:53:38) [GCC 11.3.0]"
    },
    "vscode": {
      "interpreter": {
        "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}