TensorFlow object detection with custom objects

4 min readJun 25, 2018

Quick demo of object detection by TensorFlow

We are creating a model that can identify hardware tools using by TensorFlow. In order to train the TensorFlow model, we will need to two files — custom object Model files (.pb) and object names file (.pbtxt). You can read more about these files and how to use them here.

I have laid out the steps to train our model below. These instructions apply to Ubuntu and I hope it will be helpful.

Step 1:- Install Tensorflow by following the instructions here. You will also need Python installed for this.

Step 2:- The TensorFlow project has some very useful API and scripts that can help us train our model. Go ahead and clone this repository from here.

Step 3:- In order to train our model, we need to provide it with some training data. In our case, since we want it to detect tools, we need to collect images of various tools. Make sure you collect a ton of images and you can also datasets online for this.

Step 4:- Once images are ready, we need to annotate them to describe the object in the image. To do this we can use a powerful graphical annotation tool called labellmg which generates an XML file for you.

Step 5:- TensorFlow expects training data as TFRecord file. In order to convert the XML file we obtained from labelimg, we first need to convert it to CSV using xml_to_csv.py here and then use generate_tfrecord.py here to get the TFRecord file.

While using the xml_to_csv.py script, make sure you change the path to your images folder and the output CSV folder.

image_path = os.path.join(os.getcwd(),’your image path’.format(directory))xml_df.to_csv(‘your csv folder ’.format(directory), index=None)

Here’s how you use generate_tfrecord.py to convert your CSV to TFRecord.

$ python generate_tfrecord.py --csv_input=screwdriver_data/train_labels.csv --output_path=screwdriver_data/train.record

To set the object name in TFRecord you can change the generate_tfrecord.py file in this place —

# TO-DO replace this with label map
def class_text_to_int(row_label):
 print(row_label)
 if row_label == ‘Your object name’:
 return 1
 else:
 None

You might get errors at this point if Protobuf is not compiled properly. You can follow the troubleshooting steps here.

Step 7:- The next step is to train a TensorFlow model. But instead of starting from scratch, let’s use a pre-trained model and re-config so that it can be trained to detect our custom objects, tools in our case. You can follow instructions for this here.

Step 8:- Clone the Tensorflow model repository and navigate to the research/object_detection folder and then execute the below commands in this path.

python train.py --logtostderr --train_dir=training/ --pipeline_config_path=training/ssd_mobilenet_v1_pets.config

I recommend you take a coffee break as the above process could take a long while…

All files will create in the training folder. You can visualize the training progress by using TensorBoard by executing the below command in object_detection folder—

$ tensorboard --logdir='training'

Step 9:- Finally, execute the below command in the same path which will generate the model file(.bp) in the specified output folder. You should also use the latest “Trained Checkpoint Prefix” filename which is of the format — model.ckpt-XX in the below command.

python export_inference_graph.py \
 — input_type image_tensor \
 — pipeline_config_path training/ssd_mobilenet_v1_pets.config \
 — trained_checkpoint_prefix training/model.ckpt-XX\
 — output_directory “graphfoldername”

And that is it. You have your trained model that you can use to detect the objects you want. Finally, you can play with custom object detection by TensorFlow.

.In the next blog I will write about how to use this model along with OpenCV to build an object detection solution to generate outputs like the above image.

References:

TensorFlow object detection with custom objects

Written by Bala Venkatesh