Training Tensorflow Object Detection API with custom dataset for working in Javascript and Vue.js
By: Adrià Gil
Tags: JavaScript, Artificial Intelligence, Guides And Tutorials, TensorFlow, Machine Learning

Generate the TFRecords for training

In order to create the TFRecords we will use two scripts from Dat Tran’s raccoon detector; xml_to_csv.py and generate_tfrecord.py files.

Download and place them in object_detection folder.

You can go to References section below to see from where I downloaded them.

We need now to modify xml_to_csv.py script so we can transform the created xml files to csv correctly.

# Old:

def main():

image_path = os.path.join(os.getcwd(), 'annotations')

xml_df = xml_to_csv(image_path)

xml_df.to_csv('raccoon_labels.csv', index=None)

print('Successfully converted xml to csv.') # New:

def main():

for folder in ['train', 'test']:

image_path = os.path.join(os.getcwd(), ('images/' + folder))

xml_df = xml_to_csv(image_path)

xml_df.to_csv(('images/'+folder+'_labels.csv'), index=None)

print('Successfully converted xml to csv.')

Then we can use the script opening the command line and typing:

python xml_to_csv.py

As you can observe, two files have been created in the images directory. One called test_labels.csv and another one called train_labels.csv

Note: if you get “No module named ‘pandas’” error just do a conda install pandas.

Before we can transform these csv files to TFRecords we need to change the generate_tfrecords.py script.

From:

# TO-DO replace this with label map

def class_text_to_int(row_label):

if row_label == 'raccoon':

return 1

else:

None

To:

def class_text_to_int(row_label):

if row_label == 'jagermeister bottle':

return 1

else:

None

Now the TFRecords can be generated by typing:

# train tfrecord

python generate_tfrecord.py --csv_input=images/train_labels.csv --image_dir=images/train --output_path=train.record # test tfrecord

python generate_tfrecord.py --csv_input=images/test_labels.csv --image_dir=images/test --output_path=test.record

These two command generate a train.record and a test.record file which can be used to train our object detector.

Note: if getting an error like “module tensorflow has no attribute app” is because you are using Tensorflow 2.0 so we need to change a line in generate_tfrecord.py file.

From:

# line 17

import tensorflow as tf

To:

# line 17

import tensorflow.compat.v1 as tf

Or if you prefer you can rollback to Tensorflow 1.0 instead doing:

conda remove tensorflow

conda install tensorflow=1

Configure training

Before we start training we need to create a label map and a training configuration file.

Creating a label map

The label map maps an id to name. We will put it in a folder called training located in the object_detection directory with the name labelmap.pbtxt

item {

id: 1

name: 'jagermeister bottle'

}

The id number of each item should match the id of specified item in the generate_tfrecord.py file.

Creating a training configuration

Now we need to create a training configuration file.

We are going to use faster_rcnn_inception_v2_coco model which can be downloaded from:

https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/detection_model_zoo.md

Download it uncompress the file and place it into object_detection folder.

Folder name will look like faster_rcnn_inception_v2_coco_2018_01_28

And we are going to start with a sample config file named faster_rcnn_inception_v2_pets.config which can be found in the sample folder.

You can download it from:

https://github.com/tensorflow/models/tree/master/research/object_detection/samples/configs

Keep the same name and save it into the training folder and open it with a text editor in order to change a few lines of code.

Line 9: change the number of classes to number of objects you want to detect (1 in our case).

Line 106: change fine_tune_checkpoint to the path of the model.ckpt file

fine_tune_checkpoint:

"/Users/<username>/projects/tensorflow/models/research/object_detection/faster_rcnn_inception_v2_coco_2018_01_28/model.ckpt"

Line 123: change input_path to the path of the train.record file:

input_path:

"/Users/<username>/projects/tensorflow/models/research/object_detection/train.record

Line 135: change input_path to the path of the test.records file:

input_path:

"/Users/<username>/projects/tensorflow/models/research/object_detection/test.record

Line 125-137: change label_map_path to the path of label map file:

label_map_path:

"/Users/<username>/projects/tensorflow/models/research/object_detection/training/labelmap.pbtxt

Line 130: change num_example to the number of images in your test folder.

num_examples: 10

Training model

We are going to use the train.py file which is located in the object_detection/legacy folder. We will copy it into the object_detection folder and then we will type the following into the command line:

Update: or we can use the model_main.py file in the object_detection folder instead.

python model_main.py --logtostderr --model_dir=training/ --pipeline_config_path=training/faster_rcnn_inception_v2_pets.config

If everything was setup correctly the training should begin shortly.

Note: if you are getting an error like “module tensorflow has no attribute contrib” is because you are using Tensorflow 2.0.

There are two ways of solving this:

update scripts for being used with Tensorflow 2.0 rollback to Tensorflow 1.0

The easy way is doing a rollback so do the following

# uninstall Tensorflow 2.0

conda remove tensorflow # install Tensorflow 1.X

conda install tensorflow=1

Training model progress

Once the training is running you will see every 5 minutes or so (depending on your hardware) the current loss gets logged to Tensorboard. We can open Tensorboard by opening a second command line, navigating to object_detection folder and typing:

tensorboard --logdir=training

A web page is now opened at localhost:6006

You should train the model until it reaches a satisfying loss.

The training process can be then terminated by pressing Ctrl+C.