Training Tensorflow Object Detection API with custom dataset for working in Javascript and Vue.js | by Adrià Gil | Towards Data Science

This article is a re-work of the amazing tutorial from Gilbert Tanner on how to create your own object detector with Tensorflow Object Detection API.I’m following exactly the same steps but with some differences and adding some things I’ve faced during setup and training.
I want to thank the author for the original content and of course give credit to him.
In this article we are going to use OS X + Anaconda environment so then can be easily portable to Google Cloud Platform for example if you wish...with the addition of using GPUs instead of CPUs for computing 🚀
Also I have to say that I have developed this experiment in my research time working at Manifiesto where I perform as a frontend/online developer 🤗.
Download and install Anaconda (Python 3.7 for OS X in my case) from:https://www.anaconda.com/distribution/
I will use the same version of Python that the actual Anaconda installation ships with for the new environment that we are going to create in a minute.To know which version simply run:
conda info
Then create and activate the environment:
conda create --name tf-object-detection python=3.7.4conda activate tf-object-detection
Install Tensorflow models and dependenciesChoose a folder you want to work at and create a directory named tensorflow.I will be working at the following location:
/Users/<username>/projects/tensorflow
Then just clone the Tensorflow models repository running:
git clone https://github.com/tensorflow/models
Now you need to install all the dependencies:
conda install Cythonconda install contextlib2conda install pillowconda install lxmlconda install jupyterconda install matplotlibconda install tensorflow=1
Install the COCO APIFor this go back to your projects folder and clone the Github project repository and execute the commands below:
git clone https://github.com/cocodataset/cocoapi.gitcd cocoapi/PythonAPImakecp -r pycocotools <path_to_tensorflow>/models/research
Protobuf Installation/CompilationThis is because Tensorflow Object Detection API use .proto files and these files need to be compiled into .py files in order for the Object Detection API to work properly. Protobuf can compile these files.
Protobuf can be downloaded from:https://github.com/protocolbuffers/protobuf/releasesPlace the downloaded file anywhere you want, for example in the projects folder.
After extracting you need to go to models/research and use protobuf to extract python files from the proto files in the object_detection/protos directory.For this purpose we are going to use a little script.Save it inside the research folder and name it use_protobuf.py
import osimport sysargs = sys.argvdirectory = args[1]protoc_path = args[2]for file in os.listdir(directory):    if file.endswith(".proto"):        os.system(protoc_path+" "+directory+"/"+file+" --python_out=.")
And then you can use it:
python use_protobuf.py <path_to_directory> <path_to_protoc_file># for example in our particular casepython use_protobuf.py object_detection/protos /Users/<username>/projects/protoc/bin/protoc
Add necessary environment variablesWe need to add the research and research/slim folder to our environment variables and run the setup.py file.
export PYTHONPATH=$PYTHONPATH:<PATH_TO_TF>/TensorFlow/models/researchexport PYTHONPATH=$PYTHONPATH:<PATH_TO_TF>/TensorFlow/models/research/object_detectionexport PYTHONPATH=$PYTHONPATH:<PATH_TO_TF>/TensorFlow/models/research/slim
Note: the $PYTHONPATH environment variables needs to be added every time the shell session is closed or the Anaconda environment is deactivated.
Now navigate to tensorflow/models/research and run:
python setup.py buildpython setup.py install
Testing Tensorflow Object Detection APIAfter the installation is complete we can test everything is working correctly by running the object_detection_tutorial.ipynb from the object_detection folder.
Note: is important to have in consideration that this tutorial works for Tensorflow 2.0 and you must have Tensorflow installed in your environment — if not just run conda install tensorflow=2
jupyter notebook
Then select the object_detection_tutorial.ipynb from the browser window and just follow the instructions.
You can also check everything is working by simply importing object_detection inside a python shell:
import object_detection
And if no errors then you can assume everything is working properly.
For our purpose we are going to use a bunch of images representing Jagermeister bottles.To do that as fast as possible we are going to scrap google images with a python script to automate the process.
Scripts can be downloaded from:https://github.com/hardikvasa/google-images-downloadand, for example, place the downloaded files in projects folder named as google-images-download.
Then just navigate to the folder where the python script is and execute:
python google_images_download.py --keywords "jagermeister bottle" --limit 100 --format jpg
Because the images resolution may be very different and some can be quite big we want to scale them to do the training process faster.Here is a little script:
from PIL import Imageimport osimport argparsedef rescale_images(directory, size):    for img in os.listdir(directory):        im = Image.open(directory+img)        im_resized = im.resize(size, Image.ANTIALIAS)        im_resized.save(directory+img)if __name__ == '__main__':    parser = argparse.ArgumentParser(description="Rescale images")    parser.add_argument('-d', '--directory', type=str, required=True, help='Directory containing the images')    parser.add_argument('-s', '--size', type=int, nargs=2, required=True, metavar=('width', 'height'), help='Image size')    args = parser.parse_args()    rescale_images(args.directory, args.size)
To use the script we need to save it, for example, with the name transform_image_resolution.py and then go into the command line and type:
python transform_image_resolution.py -d <image_dir>/ -s 800 600
Note: images in the image_dir folder will be overwritten so do a backup if needed.
Note: important to check the downloaded images folder in order to avoid duplicates and delete corrupted files if any.
Now we have to have to move about 80% of images into object_detection/images/train directory and the other 20% into object_detection/images/test directory.
In order to label our data, we need some kind of labeling software.In this particular case we are going to use LabelImg.
As far as we are using Anaconda we only need to follow these instructions:
# install pyqt (version 5)conda install pyqt# download LabelImg and place it in 'projects' foldergit clone https://github.com/tzutalin/labelImg.git# navigate to labelImg directorycd labelImg# execute labelImg.pymake qt5py3;./labelImg.py
Open each train and test folders and then use the “Create RectBox” button for labeling each image and click save.We are going to use “jagermeister bottle” label for each box.After saving you will see an XML file appearing in the same directory with the same name for the image we have just labeled.
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 mapdef 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 tfrecordpython generate_tfrecord.py --csv_input=images/train_labels.csv --image_dir=images/train --output_path=train.record# test tfrecordpython 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 17import tensorflow as tf
To:
# line 17import tensorflow.compat.v1 as tf
Or if you prefer you can rollback to Tensorflow 1.0 instead doing:
conda remove tensorflowconda install tensorflow=1
Before we start training we need to create a label map and a training configuration file.
Creating a label mapThe 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 configurationNow 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
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.0rollback to Tensorflow 1.0
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.0conda remove tensorflow# install Tensorflow 1.Xconda install tensorflow=1
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.
Once the model is trained we need to generate an inference graph, which can be used to run the model. For doing so we need to first find out the highest saved step number.For this, we need to navigate to the training directory and look for the model.ckpt file with the biggest index.
Then we can create the inference graph by typing the following in the command line:
python export_inference_graph.py --input_type image_tensor --pipeline_config_path training/faster_rcnn_inception_v2_pets.config --trained_checkpoint_prefix training/model.ckpt-XXXX --output_directory inference_graph
where XXXX represents the highest number.
First we need to install TensorflowJS in our environment.Usually we would use conda to install the package but unfortunately is not available in the repository sources.So we are going to use it via python pip:
pip install tensorflowjs
Once the installation finishes we can execute the following command in order to convert our saved model to something tensorfowjs understands:
tensorflowjs_converter --input_format=tf_saved_model <input_dir> <output_dir>
For example if we are in the directory /Users/<username>/projects/tensorflow/models/research/object_detection we can run something like:
tensorflowjs_converter --input_format=tf_saved_model --output_node_names='detection_boxes,detection_classes,detection_features,detection_multiclass_scores,detection_scores,num_detections,raw_detection_boxes,raw_detection_scores' --saved_model_tags=serve --output_format=tfjs_graph_model inference_graph/saved_model inference_graph/web_model
Note that the conversion options are specifically set for converting an object detection SavedModel and the output_node_names configuration have been retrieved with thesaved_model_cli command:
saved_model_cli show --dir inference_graph/saved_model --tag_set serve --signature_def serving_default#output will print something like this:#detection_boxes#detection_classes#detection_features#detection_multiclass_scores#detection_scores#num_detections#raw_detection_boxes#raw_detection_scores
Full explanation on gist in References section below.
Now that we have our model properly converted we can use it in a web application environment that runs Javascript.For making things smoother we can head to the demo application listed in References section below.The repository you are looking for is the adriagil/tfjs-vue-example
I have to stop for a moment and thank to freshsomebody for sharing this great Vue.js demo application which uses a pre-trained coco-ssd model to be used with a webcam in a browser.I’ve just adapted the code to work with my trained model consisting basically in how model serving and loading is done and how to retrieve the needed tensors to analyze data.Here the original app in case you are interested https://github.com/freshsomebody/tfjs-coco-ssd-vue-example
Just navigate there and clone or download at your preferred location.Then copy the converted model in the previous step into the root directory of the downloaded project, the web_model folder.
Then navigate to the application root directory and run:
npm install
Once finished open the App.vue file in your text editor of choice.You can see now the source code of the web application, great! 🍺
In order to load the converted model properly we need to do it via a web server.We’re going to use a Node package to do so called http-server.All we need to do is activate the web server with the command:
node_modules/http-server/bin/http-server -c1 --cors .
and the model will be accesible via http at http://localhost:8081/web_model/model.json to be loaded in the browser.
Last but not least, we need to compile or serve the frontend application to test everything with the following command:
npm run serve
In http://localhost:8080 you should now see the webcam stream and the corresponding bounding boxes wrapping the Jaggermeister bottle in red.Awesome, we are done! 👏👏👏
Despite the results of this experience are OK there are 2 important things to have in consideration.
The first one is that working in OS X is not the best choice.I would recommend you to switch to a Linux or Windows environment but only if you are going to use GPU, if not is the same piece of crap 😅A better solution (we know proper GPU cards are very expensive) is to do it on Google Cloud Plaform.They offer a very good environment and free credits for you to spend in any of your training experiments.
The second thing that sucks is that the inception model we have chosen is precise but is very slow on computer browsers and even worse in mobile.So I’m just now training another model based on MobileNet which should perform a lot better 🤞 hopefully.Also my model is not optimized and maybe that could help, but I need to do some research on how to achieve it. We’ll see.
If you liked, want more like this or have any questions please don’t hesitate to contact me...if you arrived here, you deserve it!And again many thanks to the authors that have motivated me to write this article.
That’s all folks!
[Update 24/02/2020]As I commented before, the performance for a web browser was not acceptable for a real-time video analysis so I decided to do a bit of research.
Thanks to stackoverflow community ❤️ I found many suggestions like optimizing the model or going back to the training process again with another model. I decided to try with another model because the rcnn-inception model I trained was very precise but not really thought for the application I needed.
I picked ssd_mobilenet_v2_coco this time.After a weekend training the model and almost 30k steps — I know is not necessary for development but worth to try like in a real situation— was time for testing. Results were far better than in my first attempt so now prediction rate is almost fluid with an acceptable frame rate.
So lesson learned; choose the right model before training to suit your needs.