TensorFlow object detection with custom objects

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.

Image for post
Image for post

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):
if row_label == ‘Your object name’:
return 1

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…

Image for post
Image for post
Console in command window

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'
Image for post
Image for post
Loss graph in TensorBord

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.

Image for post
Image for post
Custom object detection

.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.


Image for post
Image for post

Written by

I have a passion for understanding technology at a fundamental level and Sharing ideas and code. * Aspire to Inspire before I expire* https://balavenkatesh.com

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store