train_model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Train and Create Doordetect Model\r\n",
    "This notebook can be used to train a custom (pre-trained) network and convert it for use on an OAK device.\r\n",
    "\r\n",
    "The code is based on this tutorial: https://tensorflow-object-detection-api-tutorial.readthedocs.io/en/latest/training.html"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Install Libraries"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### numpy\r\n",
    "Used for image manupulation"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "!pip install numpy"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### pillow\r\n",
    "Dropin replacement for PIL (Python Image Library)\r\n",
    "\r\n",
    "Used form image manipulation"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "!pip install pillow"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### tensorflow\r\n",
    "Used to train the neural network"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "!pip install tensorflow"
   ],
   "outputs": [],
   "metadata": {
    "scrolled": true,
    "tags": []
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### tensorflow object detection API\r\n",
    "Used to train object detection networks"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "!git clone https://github.com/tensorflow/models.git\r\n",
    "%cd models/research\r\n",
    "!protoc object_detection/protos/*.proto --python_out=.\r\n",
    "!cp object_detection/packages/tf2/setup.py .\r\n",
    "!pip install .\r\n",
    "%cd"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "Test it"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "%cd models/research\r\n",
    "!python3 object_detection/builders/model_builder_tf2_test.py\r\n",
    "%cd"
   ],
   "outputs": [],
   "metadata": {
    "scrolled": true,
    "tags": []
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### OPTIONAL: tensorflow_gpu\r\n",
    "Used for training on GPU"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "!pip install tensorflow_gpu"
   ],
   "outputs": [],
   "metadata": {
    "scrolled": true,
    "tags": []
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Download Dataset\r\n",
    "We are going to use an open door detection datase.\r\n",
    "\r\n",
    "If you want to use another dataset you may need to make changes in later steps (especially in [Prepare the Dataset](#prepare_the_dataset))"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "!git clone https://github.com/MiguelARD/DoorDetect-Dataset.git"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Prepare the Dataset\r\n",
    "This converts the labels/annotations to a TFRecords file which will be used to train the model using the TF2 object detection API.\r\n",
    "\r\n",
    "_NOTE:_ This is specific to the format used in the DoorDetect-dataset and you may need to alter this if you want to use it with another dataset"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import tensorflow as tf\r\n",
    "from object_detection.utils import dataset_util, label_map_util\r\n",
    "from PIL import Image\r\n",
    "import os\r\n",
    "import io\r\n",
    "import random\r\n",
    "from pathlib import Path\r\n",
    "\r\n",
    "images_path = \"DoorDetect-Dataset/images\"\r\n",
    "labels_path = \"DoorDetect-Dataset/labels\"\r\n",
    "label_map_path = \"labelmap.pbtxt\"\r\n",
    "output_train = \"train.tfrecords\"\r\n",
    "output_eval = \"eval.tfrecords\"\r\n",
    "\r\n",
    "# Load labels\r\n",
    "label_map_dict = label_map_util.get_label_map_dict(label_map_path)\r\n",
    "labels = {v: k for k, v in label_map_dict.items()}\r\n",
    "\"\"\"\r\n",
    "Alternatively use hardcoded labels\r\n",
    "labels = [\r\n",
    "    b\"door\",\r\n",
    "    b\"handle\",\r\n",
    "    b\"cabinet door\",\r\n",
    "    b\"refrigerator door\",\r\n",
    "]\r\n",
    "\"\"\"\r\n",
    "\r\n",
    "\"\"\"\r\n",
    "Clamp the input number between minumum and maximum.\r\n",
    "\"\"\"\r\n",
    "def clamp(x, minimum=0, maximum=1):\r\n",
    "    return min(max(x, minimum), maximum)\r\n",
    "\r\n",
    "\"\"\"\r\n",
    "Create a single example from an image and an annotation-file\r\n",
    "\"\"\"\r\n",
    "def create_example(image_path, annotation_path):\r\n",
    "    #print(f\"Processing {image_path.name}\")\r\n",
    "    with image_path.open('rb') as image:\r\n",
    "        encoded_jpg = image.read()\r\n",
    "    encoded_jpg_io = io.BytesIO(encoded_jpg)\r\n",
    "    image = Image.open(encoded_jpg_io)\r\n",
    "\r\n",
    "    filename = image_path.name.encode(\"utf8\")\r\n",
    "    id = image_path.stem.encode(\"utf8\")\r\n",
    "    width, height = image.size\r\n",
    "    xmins = []\r\n",
    "    xmaxs = []\r\n",
    "    ymins = []\r\n",
    "    ymaxs = []\r\n",
    "    classes = []\r\n",
    "    \r\n",
    "    with annotation_path.open('r') as annotations:\r\n",
    "        for annotation in annotations:\r\n",
    "            # Each image may have 0 or more annotations each of the form\r\n",
    "            #   class center_x center_y width height\r\n",
    "            # with\r\n",
    "            #   class: the class of annotated object. Note: The classes/labels in the files start with 0. However the object detection API reserves 0 as background, so we will add 1\r\n",
    "            #   center_x, center_y: the center of the bounding box\r\n",
    "            #   width, height: the dimensions of the bounding box\r\n",
    "            # All coordinates are normalized to the image size and are between 0 and 1.\r\n",
    "            data = annotation.split(\" \")\r\n",
    "            clss = int(data[0]) + 1\r\n",
    "            cx = float(data[1])\r\n",
    "            cy = float(data[2])\r\n",
    "            w = float(data[3]) / 2\r\n",
    "            h = float(data[4]) / 2\r\n",
    "\r\n",
    "            # The TFRecords format uses min and max position as bounding boxes instead of centerpoint & dimensions.\r\n",
    "            # So we have to convert them.\r\n",
    "            # Also make sure that everything is between 0 and 1.\r\n",
    "            xmin = clamp(cx - w)\r\n",
    "            xmax = clamp(cx + w)\r\n",
    "            ymin = clamp(cy - h)\r\n",
    "            ymax = clamp(cy + h)\r\n",
    "            \r\n",
    "            classes.append(clss)\r\n",
    "            xmins.append(xmin)\r\n",
    "            xmaxs.append(xmax)\r\n",
    "            ymins.append(ymin)\r\n",
    "            ymaxs.append(ymax)\r\n",
    "\r\n",
    "            #print(f\"\\t{labels[clss]}({clss}): ({xmin}, {ymin}) / ({xmax}, {ymax})\")\r\n",
    "    \r\n",
    "    return tf.train.Example(features=tf.train.Features(feature={\r\n",
    "        'image/height': dataset_util.int64_feature(height),\r\n",
    "        'image/width': dataset_util.int64_feature(width),\r\n",
    "        'image/filename': dataset_util.bytes_feature(filename),\r\n",
    "        'image/source_id': dataset_util.bytes_feature(id),\r\n",
    "        'image/encoded': dataset_util.bytes_feature(encoded_jpg),\r\n",
    "        'image/format': dataset_util.bytes_feature(b\"jpg\"),\r\n",
    "        'image/object/bbox/xmin': dataset_util.float_list_feature(xmins),\r\n",
    "        'image/object/bbox/xmax': dataset_util.float_list_feature(xmaxs),\r\n",
    "        'image/object/bbox/ymin': dataset_util.float_list_feature(ymins),\r\n",
    "        'image/object/bbox/ymax': dataset_util.float_list_feature(ymaxs),\r\n",
    "        'image/object/class/text': dataset_util.bytes_list_feature([labels[x].encode(\"utf8\") for x in classes]),\r\n",
    "        'image/object/class/label': dataset_util.int64_list_feature(classes),\r\n",
    "    }))\r\n",
    "\r\n",
    "\r\n",
    "\"\"\"\r\n",
    "Iterate over all images and create a records file.\r\n",
    "output_path: the path where to write the TFRecords-file\r\n",
    "image_paths: a list of image files. The corresponding annotation files are automatically searched in `labels_path`\r\n",
    "\r\n",
    "TODO: Do the association of image and annotation somewhere elses\r\n",
    "\"\"\"\r\n",
    "def write_record(output_path, image_paths):\r\n",
    "    writer = tf.io.TFRecordWriter(output_path)\r\n",
    "    cnt = 0\r\n",
    "    # Iterate over all image_paths and check if there is a corresponding annotation\r\n",
    "    # If it exists, create an example and write it to the TFRecord\r\n",
    "    for image_path in image_paths:\r\n",
    "        id = image_path.stem\r\n",
    "        annotation_path = (Path(labels_path) / id).with_suffix(\".txt\")\r\n",
    "        if image_path.is_file() and annotation_path.is_file():\r\n",
    "            example = create_example(image_path, annotation_path)\r\n",
    "            writer.write(example.SerializeToString())\r\n",
    "            cnt += 1\r\n",
    "        else:\r\n",
    "            print(f\"WARNING: Image file without annotation ({id})\")\r\n",
    "\r\n",
    "    writer.close()\r\n",
    "    print(f\"Processed {cnt} file(s)\")\r\n",
    "\r\n",
    "\"\"\"\r\n",
    "Partitions a list randomly into two.\r\n",
    "percent items go into the first list and (1-percent) into the second.\r\n",
    "\"\"\"\r\n",
    "def partition(list_in, percent=0.9):\r\n",
    "    random.shuffle(list_in)\r\n",
    "    pivot = int(percent*len(list_in))\r\n",
    "    return list_in[:pivot], list_in[pivot:]\r\n",
    "\r\n",
    "\r\n",
    "image_paths = list(Path(images_path).iterdir())\r\n",
    "train_paths, eval_paths = partition(image_paths)\r\n",
    "write_record(output_train, train_paths)\r\n",
    "write_record(output_eval, eval_paths)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Download Pre-Trained model\r\n",
    "You may use another model from the [object detection model zoo](https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf2_detection_zoo.md)"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "!curl http://download.tensorflow.org/models/object_detection/tf2/20200711/ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8.tar.gz -o pretrained_model.tar.gz"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Extract the model"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "!tar -xf pretrained_model.tar.gz"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Train the Model\r\n",
    "After preparing the dataset (TFRecords-file) and downloading a pretrained model we can finally train it."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Prepare the output directory"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "!mkdir -p trained_model/ssd_mobilenet_v2_fpnlite_320x320\r\n",
    "!cp ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8/pipeline.conf trained_model/ssd_mobilenet_v2_fpnlite_320x320"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "**IMPORTANT:** Before you continue, change the pipeline.config [as shown here](https://tensorflow-object-detection-api-tutorial.readthedocs.io/en/latest/training.html)"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Start the training process\r\n",
    "We will use the `model_main_tf2.py`-script bundled with the object detection API for this.\r\n",
    "\r\n",
    "The command line options are:\r\n",
    "- `model_dir`: the directory where the trained model should be written to\r\n",
    "- `pipeline_config_path`: the path to the `pipeline.config`-file"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "!python3 models/research/object_detection/model_main_tf2.py --model_dir=trained_models/ssd_mobilenet_v2_fpnlite_320x320 --pipeline_config_path=trained_models/ssd_mobilenet_v2_fpnlite_320x320/pipeline.config"
   ],
   "outputs": [],
   "metadata": {
    "scrolled": true,
    "tags": []
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Optionally start tensorboard\r\n",
    "\r\n",
    "This can be used to monitor the training progress"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "!tensorboard --logdir trained_models/ssd_mobilenet_v2_fpnlite_320x320\r\n",
    "# Add the --bind_all option if you want to access the tensorboard from another machine. This is for example necessary if you execute this in the cloud"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "You may also want to start an evaluation process to better monitor the progress of your training.\r\n",
    "\r\n",
    "This uses the same script that is used for training. However an additional argument is added:\r\n",
    "- `checkpoint_dir`: the directory containing the checkpoints. Usually the same as the `model_dir`"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "!python3 models/research/object_detection/model_main_tf2.py --model_dir=trained_models/ssd_mobilenet_v2_fpnlite_320x320 --pipeline_config_path=trained_models/ssd_mobilenet_v2_fpnlite_320x320/pipeline.config --checkpoint_dir=trained_models/ssd_mobilenet_v2_fpnlite_320x320"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Export the model\r\n",
    "In order to use the model for inference you may want to export it (freeze it).\r\n",
    "\r\n",
    "This can be done with the `exporter_main_v2.py`-script.\r\n",
    "The following options are used:\r\n",
    "- `input_type`: The type of the input. We will keep this as `image_tensor`\r\n",
    "- `pipeline_config_path`: the path to the `pipeline.config`-file\r\n",
    "- `trained_checkpoint_dir`: the directory containing the checkpoints. Usually the same as the `model_dir`\r\n",
    "- `output_directory`: where to put the output of this script"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "!python3 models/research/object_detection/exporter_main_v2.py --input_type image_tensor --pipeline_config_path trained_models/ssd_mobilenet_v2_fpnlite_320x320/pipeline.config --trained_checkpoint_dir trained_models/ssd_mobilenet_v2_fpnlite_320x320/ --output_directory exported_models/ssd_mobilenet_v2_fpnlite_320x320"
   ],
   "outputs": [],
   "metadata": {
    "scrolled": true,
    "tags": []
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "After the above script finishes you can use the exported model for inference or you may convert it to a `.blob`-file for use on an OAK-device."
   ],
   "metadata": {}
  }
 ],
 "metadata": {
  "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}