Train a custom object detection model with Tensorflow 1

We will learn how to train an object detection model from a list of pre-trained models with the dataset you created on the picsell-IA platform.

​ β€‹πŸ‘‰ ​HERE​

If you want to follow the easy version of this guide follow the link below.

Prerequisites

We assume that you have a working python3.6+ installation on your computer.

Google-Collab
Local
Google-Collab
!git clone https://github.com/Picsell-ia/training
%cd training/
!pip install -r requirements.txt
%tensorflow_version 1.x
Local
git clone https://github.com/Picsell-ia/training
cd training
pip install -r requirements.txt

Imports

Now that you have installed what's required you can start importing the different useful modules.

import sys
sys.path.append("slim")
from picsellia import Client
import picsell_utils
import tensorflow as tf

Setting up your Picsell client

We need to connect with the Picsell platform to retrieve all of your project data. You can specify in the checkout_project method a png_dir to set the path where your images are, if it's not specified you'll need to download the images from the platform.

api_token = "your_api_token"
project_token = "your_project_token"
model_name = "your_model_name"
​
clt = Client(api_token=api_token)
clt.checkout_project(project_token=project_token)
clt.checkout_network(model_name)

You don't have your tokens ?

Once the client is initialized we can start the pre-processing of our data.

Data pre-processing

Most of the pre-processing is done by the Tensorflow object-detection API. However we do need to put our data in a TFRecord format for the API to understand it and to create a protobuf config file describing all the parameters of our model, input pipeline, evaluation and training configuration ... Hopefully everything is handled by the picsell-utils package.

Downloading data

We need the annotations and images on our machine. We also need a label map, mapping the labels names to a label ID that the Tensorflow object-detection API can comprehend. When we checked out the network the annotations were downloaded and saved and the label map was generated. We simply need to run dl_pictures() to download the images from the platform if you didn't specified png_dir when checking out the project.

clt.train_test_split()
clt.dl_pictures()

TFRecord files creation

Tensorflow needs the data in a TFRecord format, a memory efficient format storing data as a sequence of binary records. We generate two .record files, one for the training and one for the evaluation. We set the annotation type to rectangle so we can extract the bounding boxes and their label from the annotations and store them inside the records.

annotation_type = "rectangle"
picsell_utils.create_record_files(label_path=clt.label_path,
record_dir=clt.record_dir, tfExample_generator=clt.tf_vars_generator,
annotation_type=annotation_type)

Editing the configuration file

Every trainable object-detection model downloaded from the Picsell.ia hub is provided with the protobuf configuration file used to train it. We want to edit this file and set some variables specific to our project. Most of those variables are provided by the SDK, but it's still up to you to set up some of them. If you want to dwell inside the configuration file to have more control over your model, you can still open it up with after that and go change some settings like the learning rate decay strategy. But it's not the purpose of this guide. The argument incremental_or_transfer is here to specify if we should load all variables from the checkpoint, in the case you want to resume a training or if we should re-init the unfrozen variables to start the transfer learning process. In short, use incremental if you want to resume a training or transfer for a new training.

#The number of steps you want your model to be trained on
nb_steps = 1000
#Your batch size, highly depending of your hardware and model architecture
#Some models will return errors while training with a really low batch size
batch_size = 16
#The learning rate used, it can be left to None to use the previous one.
learning_rate = None
​
picsell_utils.edit_config(model_selected=clt.model_selected,
config_output_dir=clt.config_dir,
record_dir=clt.record_dir,
label_map_path=clt.label_path,
num_steps=nb_steps,
batch_size=batch_size,
learning_rate=learning_rate,
annotation_type=annotation_type,
eval_number=len(clt.eval_list),
incremental_or_transfer="transfer")

Training

Now that the input pipeline is built we can finally launch the training. To do this we can use the wrapper function train from picsell_utils , we simply specify where we will save the checkpoint and where the configuration file is located.

picsell_utils.train(ckpt_dir=clt.checkpoint_dir,
conf_dir=clt.config_dir)

Once the training is done we want to send the training logs to the Picsell.ia platform.

dict_log = picsell_utils.tfevents_to_dict(path=clt.checkpoint_dir)
clt.send_logs(dict_log)

You can send the checkpoints to the Picsell.ia platform now.

clt.send_checkpoints()

Evaluating

Next is the evaluation phase, we can launch the evaluation with the evaluate function of the picsell_utils module.

metrics = picsell_utils.evaluate(clt.metrics_dir, clt.config_dir, clt.checkpoint_dir)
clt.send_metrics(metrics)

Exporting and infering

Now that the training is done and that we checked the performance of our model through the metrics returned by the evaluation we may want to export and use this model.

Exporting the model

To export our model as a saved_model.pb we use the export_infer_graph function while specifying the right paths. We can send it to the platform to be used on the playground for live inference.

picsell_utils.export_infer_graph(ckpt_dir=clt.checkpoint_dir,
exported_model_dir=clt.exported_model_dir,
pipeline_config_path=clt.config_dir)
​
clt.send_model()

Inference

Lastly we want to use the model on some images so we can send the results to the platform and see them on the dashboard. You can either specify a list of paths of images or the record_dir attribute from the client, if using a list of paths of images you need to set from_tfrecords to False. Here we used the record file for inference

picsell_utils.infer(clt.record_dir, exported_model_dir=clt.exported_model_dir,
label_map_path=clt.label_path, results_dir=clt.results_dir, from_tfrecords=True)
clt.send_examples()