Here is one interesting use-case for Machine Learning and Tensor Flow. I am going to tell you how to identify hardware objects using TensorFlow. To train the TensorFlow model, we will need 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.
Step 1:- Install TensorFlow and Python.
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:- Inorder 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 will generate a XML file for you.
Step 5:- TensorFlow expects training data as TFRecord file. Inorder to convert the XML file we obtained from labelimg, we first need to convert it to CSV using xml_to_csv.py link and then use generate_tfrecord.py link to get the TFRecord file.
While using the xml_to_csv.py script, make sure you change 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 training folder. You can visualise the training progress by using TensorBoard by executing the below command in object_detection folder—
$ tensorbord --logdir=training/
Step 9:- Finally, execute the below command in 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:
