Training_Customised_Image_Detection_Model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/jackliu333/object_detection_using_tensorflow/blob/main/Training_Customised_Image_Detection_Model.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ORSZd8cBp4Hz"
   },
   "outputs": [],
   "source": [
    "# Codes adapted from https://github.com/nicknochnack/TFODCourse/blob/main/2.%20Training%20and%20Detection.ipynb\n",
    "# Images sampled from https://github.com/marcusklasson/GroceryStoreDataset, with bounding boxes drawn manually for each object\n",
    "# To run using Google Colab; switch runtime type to GPU for speedup\n",
    "import numpy as np\n",
    "import random\n",
    "import tensorflow as tf\n",
    "\n",
    "SEED = 123\n",
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "tf.random.set_seed(SEED)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "H3m-Bouz1Rde"
   },
   "source": [
    "# Configuration parameters\n",
    "\n",
    "Including paths and training hyper-parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "vCm7RtDn0za5",
    "outputId": "9a10857b-fff7-4cfd-c568-318c1658d087"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "print(f\"Current OS: {os.name}\") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "V6_58Cc11pRG"
   },
   "outputs": [],
   "source": [
    "CUSTOM_MODEL_NAME = 'my_ssd_mobnet' \n",
    "\n",
    "# SSD has good tradeoff between speed and accuracy; can switch to other pretrained model\n",
    "PRETRAINED_MODEL_NAME = 'ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8'\n",
    "PRETRAINED_MODEL_URL = 'http://download.tensorflow.org/models/object_detection/tf2/20200711/ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8.tar.gz'\n",
    "\n",
    "# TF official script to encode training data to tf record format\n",
    "TF_RECORD_SCRIPT_NAME = 'generate_tfrecord.py'\n",
    "\n",
    "# Mapping dictionary between label and integer id\n",
    "LABEL_MAP_NAME = 'label_map.pbtxt'\n",
    "\n",
    "TENSORFLOW = '..'\n",
    "\n",
    "# Define a list of folder paths to be created (if needed) and used later\n",
    "paths = {\n",
    "    'WORKSPACE_PATH': os.path.join('Tensorflow', 'workspace'),\n",
    "    'SCRIPTS_PATH': os.path.join('Tensorflow','scripts'),\n",
    "    'APIMODEL_PATH': os.path.join('Tensorflow','models'),\n",
    "    'IMAGE_SOURCE_PATH':os.path.join('images'),\n",
    "    # bounding box annotation\n",
    "    'ANNOTATION_PATH': os.path.join('Tensorflow', 'workspace','annotations'),\n",
    "    'IMAGE_PATH': os.path.join('Tensorflow', 'workspace','images'),\n",
    "    'MODEL_PATH': os.path.join('Tensorflow', 'workspace','models'),\n",
    "    'PRETRAINED_MODEL_PATH': os.path.join('Tensorflow', 'workspace','pre-trained-models'),\n",
    "    'CHECKPOINT_PATH': os.path.join('Tensorflow', 'workspace','models',CUSTOM_MODEL_NAME), \n",
    "    'OUTPUT_PATH': os.path.join('Tensorflow', 'workspace','models',CUSTOM_MODEL_NAME, 'export'), \n",
    "    'PROTOC_PATH':os.path.join('Tensorflow','protoc')\n",
    "}\n",
    "\n",
    "files = {\n",
    "    'PIPELINE_CONFIG':os.path.join('Tensorflow', 'workspace','models', CUSTOM_MODEL_NAME, 'pipeline.config'),\n",
    "    'TF_RECORD_SCRIPT': os.path.join(paths['SCRIPTS_PATH'], TF_RECORD_SCRIPT_NAME), \n",
    "    'LABELMAP': os.path.join(paths['ANNOTATION_PATH'], LABEL_MAP_NAME)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yfIa-vW48SJR"
   },
   "outputs": [],
   "source": [
    "# create folder paths if not exist\n",
    "for p in paths.values():\n",
    "    if not os.path.exists(p):\n",
    "        !mkdir -p {p}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aQLhQWeWGLwc"
   },
   "source": [
    "# Set up TF model training logistics\n",
    "\n",
    "### Download TF model training utility scripts from TF model zoo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "fupwxPvIFsBe",
    "outputId": "950331c2-7e3a-4626-e399-73aaf5e97a72"
   },
   "outputs": [],
   "source": [
    "if not os.path.exists(os.path.join(paths['APIMODEL_PATH'], 'research', 'objection_detection')):\n",
    "    !git clone https://github.com/tensorflow/models {paths['APIMODEL_PATH']}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yXsp4NM_HQn1"
   },
   "source": [
    "### Install TF object detection library\n",
    "\n",
    "May hit some error during installation, can be ignored, as long as it returns \"OK\" in the last line."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Linux version\n",
    "if os.name=='posix':  \n",
    "    !apt-get install protobuf-compiler\n",
    "    !cd Tensorflow/models/research && protoc object_detection/protos/*.proto --python_out=. \n",
    "    !cd Tensorflow/models/research && cp object_detection/packages/tf2/setup.py . && python -m pip install . \n",
    "\n",
    "\n",
    "#windows version\n",
    "if os.name =='nt':\n",
    "    !cd Tensorflow/models/research && protoc object_detection/protos/*.proto --python_out=.\n",
    "    !cd Tensorflow/models/research && copy object_detection\\packages\\tf2\\setup.py . && python -m pip install ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "9b_j_S-uHPzC",
    "outputId": "7179bd3e-cc0a-469d-a5a8-1d67e87113df"
   },
   "outputs": [],
   "source": [
    "# Verify installation\n",
    "VERIFICATION_SCRIPT = os.path.join(paths['APIMODEL_PATH'], 'research', 'object_detection', 'builders', 'model_builder_tf2_test.py')\n",
    "!python {VERIFICATION_SCRIPT}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DGhM6ZXbJP3b"
   },
   "outputs": [],
   "source": [
    "# test if TF object detection library could be loaded\n",
    "import object_detection"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-j-GK26pKcm2"
   },
   "source": [
    "### Download and decompress TF pretrained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#linux version\n",
    "if os.name =='posix':\n",
    "    !wget {PRETRAINED_MODEL_URL}\n",
    "    !mv {PRETRAINED_MODEL_NAME+'.tar.gz'} {paths['PRETRAINED_MODEL_PATH']}\n",
    "    !cd {paths['PRETRAINED_MODEL_PATH']} && tar -zxvf {PRETRAINED_MODEL_NAME+'.tar.gz'}\n",
    "\n",
    "#Windows version\n",
    "if os.name =='nt':\n",
    "    !wget {PRETRAINED_MODEL_URL}\n",
    "    !move {PRETRAINED_MODEL_NAME+'.tar.gz'} {paths['PRETRAINED_MODEL_PATH']}\n",
    "    !cd {paths['PRETRAINED_MODEL_PATH']} && tar -zxvf {PRETRAINED_MODEL_NAME+'.tar.gz'}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KdyE_2lbPOKH"
   },
   "source": [
    "### Create label map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xuPgTDHNK0P7"
   },
   "outputs": [],
   "source": [
    "labels = [{'name':'Apple', 'id':1},\n",
    "          {'name':'Avocado', 'id':2},\n",
    "          {'name':'Banana', 'id':3},\n",
    "          {'name':'Cabbage', 'id':4},\n",
    "          {'name':'Carrot', 'id':5},\n",
    "          {'name':'Potato', 'id':6}]\n",
    "\n",
    "with open(files['LABELMAP'], 'w') as f:\n",
    "    for label in labels:\n",
    "        f.write('item { \\n')\n",
    "        f.write('\\tname:\\'{}\\'\\n'.format(label['name']))\n",
    "        f.write('\\tid:{}\\n'.format(label['id']))\n",
    "        f.write('}\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9eJQTu74vDpA"
   },
   "source": [
    "### Download training images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "HrOfuCwavHBJ",
    "outputId": "45a233b4-cd60-4608-a404-7718febf7c29"
   },
   "outputs": [],
   "source": [
    "#Not needed\n",
    "#import shutil\n",
    "\n",
    "#if os.path.exists('object_detection_using_tensorflow'):\n",
    "#    shutil.rmtree('object_detection_using_tensorflow')\n",
    "\n",
    "#!git clone https://github.com/jackliu333/object_detection_using_tensorflow.git"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zEv7XDo7nlpK"
   },
   "source": [
    "### Split into train test folders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Create image folders\n",
    "import shutil\n",
    "\n",
    "tmp_folders = ['train', 'test']\n",
    "\n",
    "#windows version\n",
    "if os.name =='nt':\n",
    "    for i in tmp_folders:\n",
    "        if os.path.exists(os.path.join(paths['IMAGE_PATH'], i)):\n",
    "            shutil.rmtree(os.path.join(paths['IMAGE_PATH'], i))\n",
    "            !mkdir {os.path.join(paths['IMAGE_PATH'], i)}\n",
    "        else:\n",
    "            !mkdir {os.path.join(paths['IMAGE_PATH'], i)}\n",
    "#Linux\n",
    "if os.name =='posix':\n",
    "    for i in tmp_folders:\n",
    "        if os.path.exists(os.path.join(paths['IMAGE_PATH'], i)):\n",
    "            shutil.rmtree(os.path.join(paths['IMAGE_PATH'], i))\n",
    "            !mkdir -p {os.path.join(paths['IMAGE_PATH'], i)}\n",
    "        else:\n",
    "            !mkdir -p {os.path.join(paths['IMAGE_PATH'], i)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5boDfQnanvSJ"
   },
   "outputs": [],
   "source": [
    "#Copy images from source folders to destination folders\n",
    "# Create a list of image file paths from \"image_source_path + label name\"\n",
    "# This assume the images are stored in folders with the category (class) name\n",
    "import shutil\n",
    "\n",
    "for i in range(len(labels)):\n",
    "    from_path = os.path.join(paths['IMAGE_SOURCE_PATH'],labels[i]['name'])\n",
    "\n",
    "    # get unique file names\n",
    "    tmp_files = os.listdir(from_path)\n",
    "    tmp_names = []\n",
    "    tmp_file_types = []\n",
    "    for tmp_file in tmp_files:\n",
    "        tmp_name = os.path.splitext(tmp_file)[0]\n",
    "        tmp_file_type = os.path.splitext(tmp_file)[1]\n",
    "        tmp_names.append(tmp_name)\n",
    "        tmp_file_types.append(tmp_file_type)\n",
    "    tmp_names = list(set(tmp_names))\n",
    "    tmp_names = [i for i in tmp_names if i != '.DS_Store']\n",
    "    tmp_file_types = list(set(tmp_file_types))\n",
    "    tmp_file_types = [i for i in tmp_file_types if len(i) != 0]\n",
    "    # random shuffle the files\n",
    "    random.shuffle(tmp_names)\n",
    "    \n",
    "    # training and test files\n",
    "    tmp_names_train = tmp_names[0:int(len(tmp_names)*0.9)]\n",
    "    tmp_names_test = [i for i in tmp_names if i not in tmp_names_train]\n",
    "\n",
    "    # move into respective target folders\n",
    "    for tmp_name in tmp_names_train:\n",
    "        for tmp_file_type in tmp_file_types:\n",
    "            tmp_name_full = tmp_name + tmp_file_type\n",
    "            shutil.copy(os.path.join(from_path, tmp_name_full), \\\n",
    "                        os.path.join(paths['IMAGE_PATH'], \"train\"))\n",
    "\n",
    "    for tmp_name in tmp_names_test:\n",
    "        for tmp_file_type in tmp_file_types:\n",
    "            tmp_name_full = tmp_name + tmp_file_type\n",
    "            shutil.copy(os.path.join(from_path, tmp_name_full), \\\n",
    "                        os.path.join(paths['IMAGE_PATH'], \"test\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RI2Xr8VAh3eX"
   },
   "source": [
    "### Create TF Record"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7HHaNqVqh5tg"
   },
   "outputs": [],
   "source": [
    "# download conversion script\n",
    "if not os.path.exists(files['TF_RECORD_SCRIPT']):\n",
    "    !git clone https://github.com/nicknochnack/GenerateTFRecord {paths['SCRIPTS_PATH']}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create train data:\n",
    "!python {os.path.join(paths['SCRIPTS_PATH'],'generate_tfrecord.py')} -x {os.path.join(paths['IMAGE_PATH'], 'train')} -l {os.path.join(paths['ANNOTATION_PATH'], 'label_map.pbtxt')} -o {os.path.join(paths['ANNOTATION_PATH'],'train.record')}\n",
    "\n",
    "# Create test data:\n",
    "!python {os.path.join(paths['SCRIPTS_PATH'],'generate_tfrecord.py')} -x {os.path.join(paths['IMAGE_PATH'], 'test')} -l {os.path.join(paths['ANNOTATION_PATH'], 'label_map.pbtxt')} -o {os.path.join(paths['ANNOTATION_PATH'],'test.record')}\n",
    "\n",
    "\n",
    "# Create train data:\n",
    "#python generate_tfrecord.py -x [PATH_TO_IMAGES_FOLDER]/train -l [PATH_TO_ANNOTATIONS_FOLDER]/label_map.pbtxt -o [PATH_TO_ANNOTATIONS_FOLDER]/train.record\n",
    "\n",
    "# Create test data:\n",
    "#python generate_tfrecord.py -x [PATH_TO_IMAGES_FOLDER]/test -l [PATH_TO_ANNOTATIONS_FOLDER]/label_map.pbtxt -o [PATH_TO_ANNOTATIONS_FOLDER]/test.record\n",
    "\n",
    "\n",
    "# For example\n",
    "# python generate_tfrecord.py -x C:/Users/sglvladi/Documents/Tensorflow/workspace/training_demo/images/train -l C:/Users/sglvladi/Documents/Tensorflow/workspace/training_demo/annotations/label_map.pbtxt -o C:/Users/sglvladi/Documents/Tensorflow/workspace/training_demo/annotations/train.record\n",
    "# python generate_tfrecord.py -x C:/Users/sglvladi/Documents/Tensorflow/workspace/training_demo/images/test -l C:/Users/sglvladi/Documents/Tensorflow2/workspace/training_demo/annotations/label_map.pbtxt -o C:/Users/sglvladi/Documents/Tensorflow/workspace/training_demo/annotations/test.record"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "AppAbu0aiGFA",
    "outputId": "994e49bf-d1e0-433b-eb04-a4491b0e6470"
   },
   "outputs": [],
   "source": [
    "# convert to TF record format\n",
    "#!python {files['TF_RECORD_SCRIPT']} -x {os.path.join(paths['IMAGE_PATH'], 'train')} -l {files['LABELMAP']} -o {os.path.join(paths['ANNOTATION_PATH'], 'train.record')} \n",
    "#!python {files['TF_RECORD_SCRIPT']} -x {os.path.join(paths['IMAGE_PATH'], 'test')} -l {files['LABELMAP']} -o {os.path.join(paths['ANNOTATION_PATH'], 'test.record')} "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "W3XsE4cQgMp5"
   },
   "source": [
    "### Copy model configuration file to training folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rmNNSD6RQT3m"
   },
   "outputs": [],
   "source": [
    "if os.name =='posix':\n",
    "    !cp {os.path.join(paths['PRETRAINED_MODEL_PATH'], PRETRAINED_MODEL_NAME, 'pipeline.config')} {paths['CHECKPOINT_PATH']}\n",
    "    \n",
    "if os.name =='nt':\n",
    "    !copy {os.path.join(paths['PRETRAINED_MODEL_PATH'], PRETRAINED_MODEL_NAME, 'pipeline.config')} {paths['CHECKPOINT_PATH']}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1JupJmSHyUSz"
   },
   "source": [
    "### Update configuration file for transfer learning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XBWbfpEexiwZ"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from object_detection.utils import config_util\n",
    "from object_detection.protos import pipeline_pb2\n",
    "from google.protobuf import text_format\n",
    "\n",
    "# Read current configuration file\n",
    "pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()\n",
    "with tf.io.gfile.GFile(files['PIPELINE_CONFIG'], \"r\") as f:                                                                                                                                                                                                                     \n",
    "    proto_str = f.read()                                                                                                                                                                                                                                          \n",
    "    text_format.Merge(proto_str, pipeline_config)  \n",
    "\n",
    "# Update based on new labels\n",
    "pipeline_config.model.ssd.num_classes = len(labels)\n",
    "pipeline_config.train_config.batch_size = 4\n",
    "pipeline_config.train_config.fine_tune_checkpoint = os.path.join(paths['PRETRAINED_MODEL_PATH'], PRETRAINED_MODEL_NAME, 'checkpoint', 'ckpt-0')\n",
    "pipeline_config.train_config.fine_tune_checkpoint_type = \"detection\"\n",
    "pipeline_config.train_input_reader.label_map_path= files['LABELMAP']\n",
    "pipeline_config.train_input_reader.tf_record_input_reader.input_path[:] = [os.path.join(paths['ANNOTATION_PATH'], 'train.record')]\n",
    "pipeline_config.eval_input_reader[0].label_map_path = files['LABELMAP']\n",
    "pipeline_config.eval_input_reader[0].tf_record_input_reader.input_path[:] = [os.path.join(paths['ANNOTATION_PATH'], 'test.record')]\n",
    "\n",
    "# Write to configuration file\n",
    "config_text = text_format.MessageToString(pipeline_config)                                                                                                                                                                                                        \n",
    "with tf.io.gfile.GFile(files['PIPELINE_CONFIG'], \"wb\") as f:                                                                                                                                                                                                                     \n",
    "    f.write(config_text)   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9b4MBsUO0koz"
   },
   "source": [
    "# Model training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ERROR: Could not find a version that satisfies the requirement opencv-python-headless==4.1.2.30 (from versions: 3.4.10.37, 3.4.11.39, 3.4.11.41, 3.4.11.43, 3.4.11.45, 3.4.13.47, 3.4.15.55, 3.4.16.59, 3.4.17.61, 3.4.17.63, 3.4.18.65, 4.3.0.38, 4.4.0.40, 4.4.0.42, 4.4.0.44, 4.4.0.46, 4.5.1.48, 4.5.3.56, 4.5.4.58, 4.5.4.60, 4.5.5.62, 4.5.5.64, 4.6.0.66, 4.7.0.68)\n",
    "ERROR: No matching distribution found for opencv-python-headless==4.1.2.30"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7_tcn61DtpdV"
   },
   "outputs": [],
   "source": [
    "# fix opencv-python version issue \n",
    "#!pip uninstall opencv-python-headless \n",
    "#!pip install opencv-python-headless==4.1.2.30"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "tt9OuLm10o2G",
    "outputId": "6d0d4be1-b689-4ee9-b8cf-f9660efae473",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "TRAINING_SCRIPT = os.path.join(paths['APIMODEL_PATH'], 'research', 'object_detection', 'model_main_tf2.py')\n",
    "command = \"python {} --model_dir={} --pipeline_config_path={} --num_train_steps=2000\".format(TRAINING_SCRIPT, paths['CHECKPOINT_PATH'],files['PIPELINE_CONFIG'])\n",
    "!{command}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PrXjxgnp4I5t"
   },
   "source": [
    "# (Optional) Model evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0huXcblU0uYm"
   },
   "outputs": [],
   "source": [
    "# stop evaluation if needed\n",
    "#command = \"python {} --model_dir={} --pipeline_config_path={} --checkpoint_dir={}\".format(TRAINING_SCRIPT, paths['CHECKPOINT_PATH'],files['PIPELINE_CONFIG'], paths['CHECKPOINT_PATH'])\n",
    "#!{command}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IXU31dxv5hwP"
   },
   "source": [
    "# Predicting object category in image\n",
    "\n",
    "### Load trained model from checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7a5XB5pQ4RzX"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import tensorflow as tf\n",
    "from object_detection.utils import label_map_util\n",
    "from object_detection.utils import visualization_utils as viz_utils\n",
    "from object_detection.builders import model_builder\n",
    "from object_detection.utils import config_util\n",
    "\n",
    "# Load pipeline config and build a detection model\n",
    "configs = config_util.get_configs_from_pipeline_file(files['PIPELINE_CONFIG'])\n",
    "detection_model = model_builder.build(model_config=configs['model'], is_training=False)\n",
    "\n",
    "# Restore checkpoint\n",
    "ckpt = tf.compat.v2.train.Checkpoint(model=detection_model)\n",
    "ckpt.restore(os.path.join(paths['CHECKPOINT_PATH'], 'ckpt-3')).expect_partial()\n",
    "\n",
    "# @tf.function\n",
    "def detect_fn(image):\n",
    "    image, shapes = detection_model.preprocess(image)\n",
    "    prediction_dict = detection_model.predict(image, shapes)\n",
    "    detections = detection_model.postprocess(prediction_dict, shapes)\n",
    "    return detections"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8xMF_w1854yf"
   },
   "source": [
    "### Detect objects in image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 269
    },
    "id": "hCDa4sGs534v",
    "outputId": "42be4cac-9a22-4076-8980-56544357ea25"
   },
   "outputs": [],
   "source": [
    "import cv2 \n",
    "from matplotlib import pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "# Randomly select an image to be detected\n",
    "tmp_img = random.choice([file for file in os.listdir(os.path.join(paths['IMAGE_PATH'], \n",
    "                                              'test')) if file.endswith(\".jpg\")])\n",
    "IMAGE_PATH = os.path.join(paths['IMAGE_PATH'], 'test', tmp_img)\n",
    "\n",
    "category_index = label_map_util.create_category_index_from_labelmap(files['LABELMAP'])\n",
    "\n",
    "img = cv2.imread(IMAGE_PATH)\n",
    "image_np = np.array(img)\n",
    "\n",
    "input_tensor = tf.convert_to_tensor(np.expand_dims(image_np, 0), dtype=tf.float32)\n",
    "detections = detect_fn(input_tensor)\n",
    "\n",
    "num_detections = int(detections.pop('num_detections'))\n",
    "detections = {key: value[0, :num_detections].numpy()\n",
    "              for key, value in detections.items()}\n",
    "detections['num_detections'] = num_detections\n",
    "\n",
    "# detection_classes should be ints.\n",
    "detections['detection_classes'] = detections['detection_classes'].astype(np.int64)\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'],\n",
    "            detections['detection_classes']+label_id_offset,\n",
    "            detections['detection_scores'],\n",
    "            category_index,\n",
    "            use_normalized_coordinates=True,\n",
    "            max_boxes_to_draw=5,\n",
    "            min_score_thresh=.5,\n",
    "            agnostic_mode=False)\n",
    "\n",
    "plt.imshow(cv2.cvtColor(image_np_with_detections, cv2.COLOR_BGR2RGB))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(detections['detection_boxes'][0:4]*[x,y,x,y])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(detections['detection_scores'][0:4])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(category_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(detections))\n",
    "print(list(detections))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(detections['detection_classes']+1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(num_detections)\n",
    "print(len(detections['detection_scores']))\n",
    "print(len(detections['detection_boxes']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y, c = image_np.shape\n",
    "print(x,y,c)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "authorship_tag": "ABX9TyNI02+tNFR8Pq6E5CJhl8VG",
   "collapsed_sections": [
    "KdyE_2lbPOKH",
    "zEv7XDo7nlpK"
   ],
   "include_colab_link": true,
   "name": "Training Customised Image Detection Model.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}