“Roboflow_TensorFlow2_Object_Detection_ipynb”的副本.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "“Roboflow-TensorFlow2-Object-Detection.ipynb”的副本",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "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/wyhpolyu2020/ObjectDetection/blob/main/%E2%80%9CRoboflow_TensorFlow2_Object_Detection_ipynb%E2%80%9D%E7%9A%84%E5%89%AF%E6%9C%AC.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fF8ysCfYKgTP"
      },
      "source": [
        "# Introduction\n",
        "\n",
        "\n",
        "In this notebook, we implement [The TensorFlow 2 Object Detection Library](https://blog.tensorflow.org/2020/07/tensorflow-2-meets-object-detection-api.html) for training on your own dataset.\n",
        "\n",
        "We also recommend reading our blog post on [Train TensorFlow 2 Object Detection on custom data](https://blog.roboflow.com/train-a-tensorflow2-object-detection-model/) side by side.\n",
        "\n",
        "We will take the following steps to implement YOLOv4 on our custom data:\n",
        "* Install TensorFlow2 Object Detection Dependencies\n",
        "* Download Custom TensorFlow2 Object Detection Dataset\n",
        "* Write Custom TensorFlow2 Object Detection Training Configuation\n",
        "* Train Custom TensorFlow2 Object Detection Model\n",
        "* Export Custom TensorFlow2 Object Detection Weights\n",
        "* Use Trained TensorFlow2 Object Detection For Inference on Test Images\n",
        "\n",
        "When you are done you will have a custom detector that you can use. It will make inference like this:\n",
        "\n",
        "#### ![Roboflow Workmark](https://i.imgur.com/L0n564N.png)\n",
        "\n",
        "### **Reach out for support**\n",
        "\n",
        "If you run into any hurdles on your own data set or just want to share some cool results in your own domain, [reach out!](https://roboflow.ai) \n",
        "\n",
        "\n",
        "\n",
        "#### ![Roboflow Workmark](https://i.imgur.com/WHFqYSJ.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l7EOtpvlLeS0"
      },
      "source": [
        "# Install TensorFlow2 Object Detection Dependencies"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ypWGYdPlLRUN"
      },
      "source": [
        "import os\n",
        "import pathlib\n",
        "\n",
        "# Clone the tensorflow models repository if it doesn't already exist\n",
        "if \"models\" in pathlib.Path.cwd().parts:\n",
        "  while \"models\" in pathlib.Path.cwd().parts:\n",
        "    os.chdir('..')\n",
        "elif not pathlib.Path('models').exists():\n",
        "  !git clone --depth 1 https://github.com/tensorflow/models"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6QPmVBSlLTzM"
      },
      "source": [
        "# Install the Object Detection API\n",
        "%%bash\n",
        "cd models/research/\n",
        "protoc object_detection/protos/*.proto --python_out=.\n",
        "cp object_detection/packages/tf2/setup.py .\n",
        "python -m pip install ."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wHfsJ5nWLWh9"
      },
      "source": [
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import os\n",
        "import random\n",
        "import io\n",
        "import imageio\n",
        "import glob\n",
        "import scipy.misc\n",
        "import numpy as np\n",
        "from six import BytesIO\n",
        "from PIL import Image, ImageDraw, ImageFont\n",
        "from IPython.display import display, Javascript\n",
        "from IPython.display import Image as IPyImage\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "from object_detection.utils import label_map_util\n",
        "from object_detection.utils import config_util\n",
        "from object_detection.utils import visualization_utils as viz_utils\n",
        "from object_detection.utils import colab_utils\n",
        "from object_detection.builders import model_builder\n",
        "\n",
        "%matplotlib inline"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wh_HPMOqWH9z"
      },
      "source": [
        "#run model builder test\n",
        "!python /content/models/research/object_detection/builders/model_builder_tf2_test.py\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VA7Zbo3RLt3W"
      },
      "source": [
        "def load_image_into_numpy_array(path):\n",
        "  \"\"\"Load an image from file into a numpy array.\n",
        "\n",
        "  Puts image into numpy array to feed into tensorflow graph.\n",
        "  Note that by convention we put it into a numpy array with shape\n",
        "  (height, width, channels), where channels=3 for RGB.\n",
        "\n",
        "  Args:\n",
        "    path: a file path.\n",
        "\n",
        "  Returns:\n",
        "    uint8 numpy array with shape (img_height, img_width, 3)\n",
        "  \"\"\"\n",
        "  img_data = tf.io.gfile.GFile(path, 'rb').read()\n",
        "  image = Image.open(BytesIO(img_data))\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",
        "def plot_detections(image_np,\n",
        "                    boxes,\n",
        "                    classes,\n",
        "                    scores,\n",
        "                    category_index,\n",
        "                    figsize=(12, 16),\n",
        "                    image_name=None):\n",
        "  \"\"\"Wrapper function to visualize detections.\n",
        "\n",
        "  Args:\n",
        "    image_np: uint8 numpy array with shape (img_height, img_width, 3)\n",
        "    boxes: a numpy array of shape [N, 4]\n",
        "    classes: a numpy array of shape [N]. Note that class indices are 1-based,\n",
        "      and match the keys in the label map.\n",
        "    scores: a numpy array of shape [N] or None.  If scores=None, then\n",
        "      this function assumes that the boxes to be plotted are groundtruth\n",
        "      boxes and plot all boxes as black with no classes or scores.\n",
        "    category_index: a dict containing category dictionaries (each holding\n",
        "      category index `id` and category name `name`) keyed by category indices.\n",
        "    figsize: size for the figure.\n",
        "    image_name: a name for the image file.\n",
        "  \"\"\"\n",
        "  image_np_with_annotations = image_np.copy()\n",
        "  viz_utils.visualize_boxes_and_labels_on_image_array(\n",
        "      image_np_with_annotations,\n",
        "      boxes,\n",
        "      classes,\n",
        "      scores,\n",
        "      category_index,\n",
        "      use_normalized_coordinates=True,\n",
        "      min_score_thresh=0.8)\n",
        "  if image_name:\n",
        "    plt.imsave(image_name, image_np_with_annotations)\n",
        "  else:\n",
        "    plt.imshow(image_np_with_annotations)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IPbU4I7aL9Fl"
      },
      "source": [
        "# Prepare Tensorflow 2 Object Detection Training Data\n",
        "\n",
        "\n",
        "Roboflow automatically creates our TFRecord and label_map files that we need!\n",
        "\n",
        "**Generating your own TFRecords the only step you need to change for your own custom dataset.**\n",
        "\n",
        "Because we need one TFRecord file for our training data, and one TFRecord file for our test data, we'll create two separate datasets in Roboflow and generate one set of TFRecords for each.\n",
        "\n",
        "To create a dataset in Roboflow and generate TFRecords, follow [this step-by-step guide](https://blog.roboflow.ai/getting-started-with-roboflow/).\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OIREg_YwDa7-"
      },
      "source": [
        "![](https://i.imgur.com/ZwMdcbY.png)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QcHJuaurS_AO"
      },
      "source": [
        "#Downloading data from Roboflow\n",
        "#UPDATE THIS LINK - get our data from Roboflow\n",
        "%cd /content\n",
        "!curl -L \"https://app.roboflow.com/ds/LvgihkdqHO?key=u7JCIqa72b\" > roboflow.zip; unzip roboflow.zip; rm roboflow.zip\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YUd2wtfrqedy"
      },
      "source": [
        "# NOTE: Update these TFRecord names from \"cells\" and \"cells_label_map\" to your files!\n",
        "test_record_fname = '/content/valid/1.tfrecord'\n",
        "train_record_fname = '/content/train/1.tfrecord'\n",
        "label_map_pbtxt_fname = '/content/train/1_label_map.pbtxt'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I2MAcgJ53STW"
      },
      "source": [
        "# Configure Custom TensorFlow2 Object Detection Training Configuration\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "> In this section you can specify any model in the [TF2 OD model zoo](https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf2_detection_zoo.md) and set up your training configuration.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gN0EUEa3e5Un"
      },
      "source": [
        "##change chosen model to deploy different models available in the TF2 object detection zoo\n",
        "MODELS_CONFIG = {\n",
        "    'efficientdet-d0': {\n",
        "        'model_name': 'efficientdet_d0_coco17_tpu-32',\n",
        "        'base_pipeline_file': 'ssd_efficientdet_d0_512x512_coco17_tpu-8.config',\n",
        "        'pretrained_checkpoint': 'efficientdet_d0_coco17_tpu-32.tar.gz',\n",
        "        'batch_size': 16\n",
        "    },\n",
        "    'efficientdet-d1': {\n",
        "        'model_name': 'efficientdet_d1_coco17_tpu-32',\n",
        "        'base_pipeline_file': 'ssd_efficientdet_d1_640x640_coco17_tpu-8.config',\n",
        "        'pretrained_checkpoint': 'efficientdet_d1_coco17_tpu-32.tar.gz',\n",
        "        'batch_size': 16\n",
        "    },\n",
        "    'efficientdet-d2': {\n",
        "        'model_name': 'efficientdet_d2_coco17_tpu-32',\n",
        "        'base_pipeline_file': 'ssd_efficientdet_d2_768x768_coco17_tpu-8.config',\n",
        "        'pretrained_checkpoint': 'efficientdet_d2_coco17_tpu-32.tar.gz',\n",
        "        'batch_size': 16\n",
        "    },\n",
        "        'efficientdet-d3': {\n",
        "        'model_name': 'efficientdet_d3_coco17_tpu-32',\n",
        "        'base_pipeline_file': 'ssd_efficientdet_d3_896x896_coco17_tpu-32.config',\n",
        "        'pretrained_checkpoint': 'efficientdet_d3_coco17_tpu-32.tar.gz',\n",
        "        'batch_size': 16\n",
        "    }\n",
        "}\n",
        "\n",
        "#in this tutorial we implement the lightweight, smallest state of the art efficientdet model\n",
        "#if you want to scale up tot larger efficientdet models you will likely need more compute!\n",
        "chosen_model = 'efficientdet-d0'\n",
        "\n",
        "num_steps = 4000 #The more steps, the longer the training. Increase if your loss function is still decreasing and validation metrics are increasing. \n",
        "num_eval_steps = 500 #Perform evaluation after so many steps\n",
        "\n",
        "model_name = MODELS_CONFIG[chosen_model]['model_name']\n",
        "pretrained_checkpoint = MODELS_CONFIG[chosen_model]['pretrained_checkpoint']\n",
        "base_pipeline_file = MODELS_CONFIG[chosen_model]['base_pipeline_file']\n",
        "batch_size = MODELS_CONFIG[chosen_model]['batch_size'] #if you can fit a large batch in memory, it may speed up your training"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kG4TmJUVrYQ7"
      },
      "source": [
        "#download pretrained weights\n",
        "%mkdir /content/models/research/deploy/\n",
        "%cd /content/models/research/deploy/\n",
        "import tarfile\n",
        "download_tar = 'http://download.tensorflow.org/models/object_detection/tf2/20200711/' + pretrained_checkpoint\n",
        "\n",
        "!wget {download_tar}\n",
        "tar = tarfile.open(pretrained_checkpoint)\n",
        "tar.extractall()\n",
        "tar.close()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c-nqYZtdtsgG"
      },
      "source": [
        "#download base training configuration file\n",
        "%cd /content/models/research/deploy\n",
        "download_config = 'https://raw.githubusercontent.com/tensorflow/models/master/research/object_detection/configs/tf2/' + base_pipeline_file\n",
        "!wget {download_config}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b_ki9jOqxn7V"
      },
      "source": [
        "#prepare\n",
        "pipeline_fname = '/content/models/research/deploy/' + base_pipeline_file\n",
        "fine_tune_checkpoint = '/content/models/research/deploy/' + model_name + '/checkpoint/ckpt-0'\n",
        "\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())\n",
        "num_classes = get_num_classes(label_map_pbtxt_fname)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5eA5ht3_yukT"
      },
      "source": [
        "#write custom configuration file by slotting our dataset, model checkpoint, and training parameters into the base pipeline file\n",
        "\n",
        "import re\n",
        "\n",
        "%cd /content/models/research/deploy\n",
        "print('writing custom configuration file')\n",
        "\n",
        "with open(pipeline_fname) as f:\n",
        "    s = f.read()\n",
        "with open('pipeline_file.config', '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: \".*?)(PATH_TO_BE_CONFIGURED/train)(.*?\")', 'input_path: \"{}\"'.format(train_record_fname), s)\n",
        "    s = re.sub(\n",
        "        '(input_path: \".*?)(PATH_TO_BE_CONFIGURED/val)(.*?\")', '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",
        "    \n",
        "    #fine-tune checkpoint type\n",
        "    s = re.sub(\n",
        "        'fine_tune_checkpoint_type: \"classification\"', 'fine_tune_checkpoint_type: \"{}\"'.format('detection'), s)\n",
        "        \n",
        "    f.write(s)\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HEsOLOMHzBqF"
      },
      "source": [
        "%cat /content/models/research/deploy/pipeline_file.config"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GMlaN3rs3zLe"
      },
      "source": [
        "pipeline_file = '/content/models/research/deploy/pipeline_file.config'\n",
        "model_dir = '/content/training/'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XxPj_QV43qD5"
      },
      "source": [
        "# Train Custom TF2 Object Detector\n",
        "\n",
        "* pipeline_file: defined above in writing custom training configuration\n",
        "* model_dir: the location tensorboard logs and saved model checkpoints will save to\n",
        "* num_train_steps: how long to train for\n",
        "* num_eval_steps: perform eval on validation set after this many steps\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tQTfZChVzzpZ"
      },
      "source": [
        "!python /content/models/research/object_detection/model_main_tf2.py \\\n",
        "    --pipeline_config_path={pipeline_file} \\\n",
        "    --model_dir={model_dir} \\\n",
        "    --alsologtostderr \\\n",
        "    --num_train_steps={num_steps} \\\n",
        "    --sample_1_of_n_eval_examples=1 \\\n",
        "    --num_eval_steps={num_eval_steps}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9KNv1N_hUibE"
      },
      "source": [
        "#run model evaluation to obtain performance metrics\n",
        "#!python /content/models/research/object_detection/model_main_tf2.py \\\n",
        "    #--pipeline_config_path={pipeline_file} \\\n",
        "    #--model_dir={model_dir} \\\n",
        "    #--checkpoint_dir={model_dir} \\\n",
        "#Not yet implemented for EfficientDet"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TI9iCCxoNlAL"
      },
      "source": [
        "%load_ext tensorboard\n",
        "%tensorboard --logdir '/content/training/train'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4Vk2146Ogil3"
      },
      "source": [
        "## Exporting a Trained Inference Graph\n",
        "Still to come for TF2 models, we will be updating this Colab notebook accordingly as the functionality is added. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vqaZ4v-vIuDl"
      },
      "source": [
        "#see where our model saved weights\n",
        "%ls '/content/training/'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YnSEZIzl4M10"
      },
      "source": [
        "#run conversion script\n",
        "import re\n",
        "import numpy as np\n",
        "\n",
        "output_directory = '/content/fine_tuned_model'\n",
        "\n",
        "#place the model weights you would like to export here\n",
        "last_model_path = '/content/training/'\n",
        "print(last_model_path)\n",
        "!python /content/models/research/object_detection/exporter_main_v2.py \\\n",
        "    --trained_checkpoint_dir {last_model_path} \\\n",
        "    --output_directory {output_directory} \\\n",
        "    --pipeline_config_path {pipeline_file}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TsE_uVjlsz3u"
      },
      "source": [
        "%ls '/content/fine_tuned_model/saved_model/'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7Vz2vJeCCyZR"
      },
      "source": [
        "# Run Inference on Test Images with Custom TensorFlow2 Object Detector"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kcR4PWC3KBau",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "4baf4450-6d75-4e10-c3e0-b1e484f9eff7"
      },
      "source": [
        "#downloading test images from Roboflow\n",
        "#export dataset above with format COCO JSON\n",
        "#or import your test images via other means. \n",
        "%mkdir /content/test/\n",
        "%cd /content/test/\n",
        "!curl -L \"https://app.roboflow.com/ds/ALFAIlaNzD?key=rG5blSjKSA\" > roboflow.zip; unzip roboflow.zip; rm roboflow.zip"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "mkdir: cannot create directory ‘/content/test/’: File exists\n",
            "/content/test\n",
            "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
            "                                 Dload  Upload   Total   Spent    Left  Speed\n",
            "100   887  100   887    0     0   1174      0 --:--:-- --:--:-- --:--:--  1174\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xxtm1NutE5vK"
      },
      "source": [
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import io\n",
        "import scipy.misc\n",
        "import numpy as np\n",
        "from six import BytesIO\n",
        "from PIL import Image, ImageDraw, ImageFont\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "from object_detection.utils import label_map_util\n",
        "from object_detection.utils import config_util\n",
        "from object_detection.utils import visualization_utils as viz_utils\n",
        "from object_detection.builders import model_builder\n",
        "\n",
        "%matplotlib inline"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qs1HJnEhyevJ"
      },
      "source": [
        "def load_image_into_numpy_array(path):\n",
        "  \"\"\"Load an image from file into a numpy array.\n",
        "\n",
        "  Puts image into numpy array to feed into tensorflow graph.\n",
        "  Note that by convention we put it into a numpy array with shape\n",
        "  (height, width, channels), where channels=3 for RGB.\n",
        "\n",
        "  Args:\n",
        "    path: the file path to the image\n",
        "\n",
        "  Returns:\n",
        "    uint8 numpy array with shape (img_height, img_width, 3)\n",
        "  \"\"\"\n",
        "  img_data = tf.io.gfile.GFile(path, 'rb').read()\n",
        "  image = Image.open(BytesIO(img_data))\n",
        "  (im_width, im_height) = image.size\n",
        "  return np.array(image.getdata()).reshape(\n",
        "      (im_height, im_width, 3)).astype(np.uint8)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0f6DTolSDfXs"
      },
      "source": [
        "%ls '/content/training/'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gFY75DfTDHaU"
      },
      "source": [
        "import pathlib\n",
        "\n",
        "filenames = list(pathlib.Path('/content/training/').glob('*.index'))\n",
        "\n",
        "filenames.sort()\n",
        "print(filenames)\n",
        "\n",
        "#recover our saved model\n",
        "pipeline_config = pipeline_file\n",
        "#generally you want to put the last ckpt from training in here\n",
        "model_dir = str(filenames[-1]).replace('.index','')\n",
        "configs = config_util.get_configs_from_pipeline_file(pipeline_config)\n",
        "model_config = configs['model']\n",
        "detection_model = model_builder.build(\n",
        "      model_config=model_config, is_training=False)\n",
        "\n",
        "# Restore checkpoint\n",
        "ckpt = tf.compat.v2.train.Checkpoint(\n",
        "      model=detection_model)\n",
        "ckpt.restore(os.path.join(str(filenames[-1]).replace('.index','')))\n",
        "\n",
        "\n",
        "def get_model_detection_function(model):\n",
        "  \"\"\"Get a tf.function for detection.\"\"\"\n",
        "\n",
        "  @tf.function\n",
        "  def detect_fn(image):\n",
        "    \"\"\"Detect objects in image.\"\"\"\n",
        "\n",
        "    image, shapes = model.preprocess(image)\n",
        "    prediction_dict = model.predict(image, shapes)\n",
        "    detections = model.postprocess(prediction_dict, shapes)\n",
        "\n",
        "    return detections, prediction_dict, tf.reshape(shapes, [-1])\n",
        "\n",
        "  return detect_fn\n",
        "\n",
        "detect_fn = get_model_detection_function(detection_model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-Ycfl7rnDT1D"
      },
      "source": [
        "#map labels for inference decoding\n",
        "label_map_path = configs['eval_input_config'].label_map_path\n",
        "label_map = label_map_util.load_labelmap(label_map_path)\n",
        "categories = label_map_util.convert_label_map_to_categories(\n",
        "    label_map,\n",
        "    max_num_classes=label_map_util.get_max_label_map_index(label_map),\n",
        "    use_display_name=True)\n",
        "category_index = label_map_util.create_category_index(categories)\n",
        "label_map_dict = label_map_util.get_label_map_dict(label_map, use_display_name=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wN1BzORoIzV4"
      },
      "source": [
        "#run detector on test image\n",
        "#it takes a little longer on the first run and then runs at normal speed. \n",
        "import random\n",
        "\n",
        "TEST_IMAGE_PATHS = glob.glob('/content/test/test/*.jpg')\n",
        "image_path = random.choice(TEST_IMAGE_PATHS)\n",
        "image_np = load_image_into_numpy_array(image_path)\n",
        "\n",
        "# Things to try:\n",
        "# Flip horizontally\n",
        "# image_np = np.fliplr(image_np).copy()\n",
        "\n",
        "# Convert image to grayscale\n",
        "# image_np = np.tile(\n",
        "#     np.mean(image_np, 2, keepdims=True), (1, 1, 3)).astype(np.uint8)\n",
        "\n",
        "input_tensor = tf.convert_to_tensor(\n",
        "    np.expand_dims(image_np, 0), dtype=tf.float32)\n",
        "detections, predictions_dict, shapes = detect_fn(input_tensor)\n",
        "\n",
        "label_id_offset = 1\n",
        "image_np_with_detections = image_np.copy()\n",
        "\n",
        "viz_utils.visualize_boxes_and_labels_on_image_array(\n",
        "      image_np_with_detections,\n",
        "      detections['detection_boxes'][0].numpy(),\n",
        "      (detections['detection_classes'][0].numpy() + label_id_offset).astype(int),\n",
        "      detections['detection_scores'][0].numpy(),\n",
        "      category_index,\n",
        "      use_normalized_coordinates=True,\n",
        "      max_boxes_to_draw=200,\n",
        "      min_score_thresh=.5,\n",
        "      agnostic_mode=False,\n",
        ")\n",
        "\n",
        "plt.figure(figsize=(12,16))\n",
        "plt.imshow(image_np_with_detections)\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zQ-N94cKB82o"
      },
      "source": [
        "## Congrats!\n",
        "\n",
        "Hope you enjoyed this!\n",
        "\n",
        "--Team [Roboflow](https://roboflow.ai)\n"
      ]
    }
  ]
}