Custom Object Detection using TensorFlow from Scratch | by Khush Patel | Towards Data Science

In this tutorial, we’re going to get our hands dirty and train our own dog (corgi) detector using a pre-trained SSD MobileNet V2 model.
Instead of training your own model from scratch, you can build on existing models and fine-tune them for your own purpose without requiring as much computing power.
Install Tensorflow using the following command:
$ pip install tensorflow
If you have a GPU that you can use with Tensorflow:
$ pip install tensorflow-gpu
$ pip install pillow Cython lxml jupyter matplotlib
Install protobuf using Homebrew (you can learn more about Homebrew here)
$ brew install protobuf
For protobuf installation on other OS, follow the instructions here.
In this tutorial, we’re going to use resources in the Tensorflow models repository. Since it does not come with the Tensorflow installation, we need to clone it from their Github repo:
First change into the Tensorflow directory:
# For example: ~/anaconda/envs/<your_env_name>/lib/python3.6/site-packages/tensorflow$ cd <path_to_your_tensorflow_installation>
Clone the Tensorflow models repository:
$ git clone https://github.com/tensorflow/models.git
From this point on, this directory will be referred to as the modelsdirectory
Every time you start a new terminal window to work with the pre-trained models, it is important to compile Protobuf and change your PYTHONPATH.
Run the following from your terminal:
$ cd <path_to_your_tensorflow_installation>/models/research/$ protoc object_detection/protos/*.proto --python_out=.$ export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim
Run a quick test to confirm that the Object Detection API is working properly:
$ python object_detection/builders/model_builder_test.py
If the result looks like the following, you’re ready to proceed to the next steps!
...............----------------------------------------------------------------------Ran 15 tests in 0.123sOK
To make this tutorial easier to follow along, create the following folder structure within the models directory you just cloned:
models     ├── annotations    |   └── xmls        ├── images    ├── checkpoints    ├── tf_record    ├── research    ...
These folders will be used to store the required components for our model as we proceed.
You can collect data in either images or video format. Here I mentioned both ways to collect data.
Data preparation is the most important part of training your own model. Since we’re going to train a corgi detector, we must collect pictures of corgis! About 200 of them would be sufficient.
I recommend using google-images-download to download images. It searches Google Images and then downloads images based on the inputs you provided. In the inputs, you can specify search parameters such as keywords, number of images, image format, image size, and usage rights.
Since we’re downloading more than 100 images at a time, we need a chromedriver in the models directory (download here). Once you have the chromedriver ready, you could use this sample command to download images. Make sure all your images are in the jpg format:
# From the models directory$ googleimagesdownload --keywords 'welsh corgi dog' \--limit 200 \--size medium \--chromedriver ./chromedriver \--format jpg
After downloading, save all images to models/images/. To make subsequent processes easier, let's rename the images as numbers (e.g. 1.jpg, 2.jpg) by running the following script:
import ospath = 'models/images/'counter = 1for f in os.listdir(path):    suffix = f.split('.')[-1]    if suffix == 'jpg' or suffix == 'png':        new = '{}.{}'.format(str(counter), suffix)        os.rename(path + f, path + new)        counter = int(counter) + 1
medium.com
Once you’ve collected all the images you need, you need to label them manually. There are many packages that serve this purpose. labelImg is a popular choice.
labelImg provides a user-friendly GUI. Plus, it saves label files (.xml) in the popular Pascal VOC format. If you want these images for training YOLO(You only Look Once) then use YOLO. Just set the current directory and Save directory as per our structure.
Here's what a labelled image looks like in labelImg:
Double check that every image has a corresponding .xml file and save them in models/annotations/xmls/.
For a large number of annotations, you can use different shortcuts mentioned below:
Ctrl + u - Load all of the images from a directoryCtrl + r - Change the default annotation target dirCtrl + s - Savew - Create a rect boxd - Next imagea - Previous imagedel - Delete the selected rect boxCtrl++ - Zoom inCtrl-- - Zoom outCtrl + d - Copy the current label and rect boxSpace - Flag the current image as verified↑→↓←Keyboard arrows to move selected rect box
Classes need to be listed in the label map. Since we’re only detecting corgis, the label map should contain only one item like the following:
item {    id: 1    name: 'corgi'}
Note that id must start from 1, because 0 is a reserved id.
Save this file as label_map.pbtxt in models/annotations/
trainval.txt is a list of image names without file extensions. Since we have sequential numbers for image names, the list should look like this:
123...198199200
Save this file as trainval.txt in models/annotations/
You can use this link to create XML files to CSV. We have all images and their bounding boxes are in XML format. Also all image has separate XML file so using this code we are creating a CSV file which contains all the XML files and their bounding box co-ordinates to single CSV file which is input for creating TFrecords.
TFRecord is an important data format designed for Tensorflow. (Read more about it here). Before you can train your custom object detector, you must convert your data into the TFRecord format.
Since we need to train as well as validate our model, the data set will be split into training (train.record) and validation sets (val.record). The purpose of training set is straight forward - it is the set of examples the model learns from. The validation set is a set of examples used DURING TRAINING to iteratively assess model accuracy.
We’re going to use create_tf_record.py to convert our data set into train.record and val.record. Download here and save it to models/research/object_detection/dataset_tools/.
Just change the label name in if row_label == ‘Label1’: as per your classifications.
This script is preconfigured to do 70–30 train-val split. Execute it by running:
# From the models directory$ python research/object_detection/dataset_tools/create_tf_record.py
If the script is executed successfully, train.record and val.record should appear in your models/research/ directory. Move them into the models/tf_record/ directory.
There are many pre-trained object detection models available in the model zoo. In order to train them using our custom data set, the models need to be restored in Tensorflow using their checkpoints (.ckpt files), which are records of previous model states.
For this tutorial, we’re going to download ssd_mobilenet_v2_coco here and save its model checkpoint files (model.ckpt.meta, model.ckpt.index, model.ckpt.data-00000-of-00001) to our models/checkpoints/ directory.
Each of the pretrained models has a config file that contains details about the model. To detect our custom class, the config file needs to be modified accordingly.
The config files are included in the models directory you cloned in the very beginning. You can find them in:
models/research/object_detection/samples/configs
In our case, we’ll modify the config file for ssd_mobilenet_v2_coco. Make a copy of it first and save it in the models/ directory.
Here are the items we need to change:
Since we’re only trying to detect corgis, change num_classes to 1fine_tune_checkpoint tells the model which checkpoint file to use. Set this to checkpoints/model.ckptThe model also needs to know where the TFRecord files and label maps are for both training and validation sets. Since our train.record and val.record are saved in tf_record folder, our config should reflect that:
Since we’re only trying to detect corgis, change num_classes to 1
fine_tune_checkpoint tells the model which checkpoint file to use. Set this to checkpoints/model.ckpt
The model also needs to know where the TFRecord files and label maps are for both training and validation sets. Since our train.record and val.record are saved in tf_record folder, our config should reflect that:
train_input_reader: {  tf_record_input_reader {    input_path: "tf_record/train.record"  }  label_map_path: "annotations/label_map.pbtxt"}eval_input_reader: {  tf_record_input_reader {    input_path: "tf_record/val.record"  }  label_map_path: "annotations/label_map.pbtxt"  shuffle: false  num_readers: 1}
At this point, your models directory should look like this:
models     ├── annotations    |   ├── label_map.pbtxt    |   ├── trainval.txt    |   └── xmls    |       ├── 1.xml    |       ├── 2.xml    |       ├── ...    |    ├── images    |   ├── 1.jpg    |   ├── 2.jpg    |   ├── ...        |    ├── checkpoints    |   ├── model.ckpt.data-00000-of-00001    |   ├── model.ckpt.index    |   └── model.ckpt.meta    |    ├── tf_record    |   ├── train.record    |   └── val.record    |    ├── research    |   ├── ...    ...
If you have successfully completed all previous steps, you’re ready to start training!
Follow the steps below:
# Change into the models directory$ cd tensorflow/models# Make directory for storing training progress$ mkdir train# Make directory for storing validation results$ mkdir eval# Begin training$ python research/object_detection/train.py \    --logtostderr \    --train_dir=train \    --pipeline_config_path=ssd_mobilenet_v2_coco.config
Training time varies depending on the computing power of your machine.
Evaluation can be run in parallel with training. The eval.py script checks the train directory for progress and evaluate the model based on the most recent checkpoint.
# From the models directory$ python research/object_detection/eval.py \    --logtostderr \    --pipeline_config_path=ssd_mobilenet_v2_coco.config \    --checkpoint_dir=train \    --eval_dir=eval
You can visualize model training progress using Tensorboard:
# From the models directory$ tensorboard --logdir=./
Based on the graphs output by Tensorboard, you may decide when you want to stop training. Usually, you may stop the process when the loss function is tapering off and no longer decreasing by a significant amount. In my case, I stopped at step 3258.
Once you finish training your model, you can export your model to be used for inference. If you’ve been following the folder structure, use the following command:
# From the models directory$ mkdir fine_tuned_model$ python research/object_detection/export_inference_graph.py \    --input_type image_tensor \    --pipeline_config_path ssd_mobilenet_v2_coco.config \    --trained_checkpoint_prefix  train/model.ckpt-<the_highest_checkpoint_number> \    --output_directory fine_tuned_model
Now that you have a model, you can use it to detect corgis in pictures and videos! For the purpose of demonstration, we’re going to detect corgis in an image. Before you proceed, pick an image you want to test the model with.
The models directory came with a notebook file (.ipynb) that we can use to get inference with a few tweaks. It is located at models/research/object_detection/object_detection_tutorial.ipynb. Follow the steps below to tweak the notebook:
MODEL_NAME = 'ssd_mobilenet_v2_coco_2018_03_29'PATH_TO_CKPT = 'path/to/your/frozen_inference_graph.pb'PATH_TO_LABELS = 'models/annotations/label_map.pbtxt'NUM_CLASSES = 1Comment out cell #5 completely (just below Download Model)Since we’re only testing on one image, comment out PATH_TO_TEST_IMAGES_DIR and TEST_IMAGE_PATHS in cell #9 (just below Detection)In cell #11 (the last cell), remove the for-loop, unindent its content, and add path to your test image:
MODEL_NAME = 'ssd_mobilenet_v2_coco_2018_03_29'
PATH_TO_CKPT = 'path/to/your/frozen_inference_graph.pb'
PATH_TO_LABELS = 'models/annotations/label_map.pbtxt'
NUM_CLASSES = 1
Comment out cell #5 completely (just below Download Model)
Since we’re only testing on one image, comment out PATH_TO_TEST_IMAGES_DIR and TEST_IMAGE_PATHS in cell #9 (just below Detection)
In cell #11 (the last cell), remove the for-loop, unindent its content, and add path to your test image:
imagepath = 'path/to/image_you_want_to_test.jpg
After following through the steps, run the notebook and you should see the corgi in your test image highlighted by a bounding box!
There you have your custom corgi detector!
Tensorflow Object Detection Model Documentation
Do visit my Website: http://www.khushpatel.com