TensorFlow 2 — Object Detection on Custom Dataset with Object Detection API

Introduction

Computer vision is a very interesting subject in the area of image analytics and a lot of use cases exist in this field. These are broadly classified as image classification, image object detection, image segmentation and image generation.

In this blog we will stick to image object detection. This topic has been gaining a lot of eye balls for commercial as well as research use cases as it deals around detecting objects in an image. For example detecting vehicles, animals, mobile phones, aircraft and any other object that you may think of. Object detection models are mostly used in self driving cars, intrusion systems, fire detection, visual inspection or even retail store analytics.

See image below:

About the problem and Dataset:

We will be using Tensorflow’s Object Detection API on TF 2 to fine-tune existing Object Detection models for a custom labelled dataset using Google Colab. There are a couple of blogs on internet today for this but they are all on TF 1.X

My special credits go to Dat Tran on all his effort for creating the Raccoon dataset and labeling all images, the dataset is available on Github — https://github.com/datitran/raccoon_dataset.

The dataset contains 200 images of raccoons which were manually box-bound labelled using the tool LabelImg to prepare the dataset in the required format.

Object Detection approach:

The object detection workflow requires the below steps:

1. Google Colab running on GPU TF2
2. Collecting the images to train and validate the Object Detection model.
3. Labeling the dataset using a tool like LabelImg.
4. Preparing a TFRecord file for ingesting in object detection API.
5. Installing the Tensorflow Object Detection API.
6. Creating the OD Config file.
7. Running the Object detection training and eval job.
8. Exporting the model.

Let’s start !!!!

1. Start the Google Colab GPU version so that we have enough horse power required to train the model.
2. Next we would want to clone the required raccoon dataset into Google Colab from the github: https://github.com/datitran/raccoon_dataset

!git clone https://github.com/datitran/raccoon_dataset.git

3. Next we would like to clone the Tensorflow Object Detection API so that we can use that in our workflow

!git clone https://github.com/tensorflow/models.git

4. There are some installation steps involved around Tensorflow 2 which can be found here, but I’ll make it easy for you:

%cd /content/models/research/
!protoc object_detection/protos/*.proto --python_out=.
# Install TensorFlow Object Detection API.
!cp object_detection/packages/tf2/setup.py .
!python -m pip install .

5. The Raccoon dataset contains two files train_labels.csv and test_labels.csv which need to be converted into TFRecord format so that it can be fed into Tensorflow’s object Detection API.

I used the below function to convert the CSV file into a TFRecord file.

Use below code to convert the CSV file into a TFRecord file:

6. Since we are focused on to quickly train the model using fine tuning approach, we are going to use the pretrained TF2 MobileNet V2 model as the feature extractor in the SSD MobileNet V2 Object Detection model. So the next logical step is to download and untar the pretrained TF2 MobileNet V2 model.

!wget http://download.tensorflow.org/models/object_detection/classification/tf2/20200710/mobilenet_v2.tar.gz
!tar -xvf mobilenet_v2.tar.gz

7. In order to run the object detection model, we need two files. A config file which contains the configuration used to train the model and a PBTXT file which contains the mapping of labels. Since we only have one label in our dataset, the PBTXT file looks something like this:

item {
  id: 1
  display_name: "raccoon"
}

8. On the other hand, I created the config file for training pipeline as below:

9. This is almost the last step, we are done with all the configuration steps. Lets start the training:

10. Monitoring of the process can be done via Tensorboard logs and we let the model train for a 1000 steps, which should take not more than 20 mins on Google Colab. After the training is complete, we can start an Evaluation job to see the results on Tensorboard.

You can see on the right hand side, the model is doing quite well on test images trained for just a 1000 steps.

11. You can export the trained checkpoints as a deployable model with a simple command. The model can be deployed either on Google AI platform or AWS Sagemaker.

I hope you enjoyed the article and going to build something new today !!!

Edit:

After I published this, I saw a few more blogs on this topic were written. One of them which is really worth reading is at https://bit.ly/3naBnQG