Training a model for custom object detection (TF 2.x)

Training a model for custom object detection (TF 2.x)

In this blog, I will be training a machine learning model for custom object detection using TensorFlow 2.x in Google Colab. Following is the roadmap for it.

Roadmap

  • Collect the dataset of images and label them to get their XML files.
  • Install the TensorFlow Object Detection API.
  • Generate the TFRecord files required for training. (need generate_tfrecord.py script and CSV files for this)
  • Edit the model pipeline config file and download the pre-trained model checkpoint.
  • Train and evaluate the model.

Here, I am training a model for custom object detection (mask). This is done in 16 steps mentioned in the section below:

  1. Import Libraries
  2. Create customTF2, training, and data folders in your google drive
  3. Create and upload your image files and XML files
  4. Upload the generate_tfrecord.py file to the customTF2 folder in your drive
  5. Mount drive and link your folder
  6. Clone the TensorFlow models git repository & Install TensorFlow Object Detection API
  7. Test the model builder
  8. Navigate to /mydrive/customTF2/data/ and Unzip the images.zip and annotations.zip files into the data folder
  9. Create test_labels & train_labels
  10. Create_CSV and “label_map.pbtxt” files
  11. Create ‘train.record’ & ‘test.record’ files
  12. Download pre-trained model checkpoint
  13. Get the model pipeline config file, make changes to it, and put it inside the data folder
  14. Load Tensorboard
  15. Train the model
  16. Test your trained model

HOW TO BEGIN?

  • Open my Colab notebook on your browser.
  • Click on File in the menu bar and click on Save a copy in drive. This will open a copy of my Colab notebook on your browser which you can now use.
  • Next, once you have opened the copy of my notebook and are connected to the Google Colab VM, click on Runtime in the menu bar and click on Change runtime type. Select GPU and click on save.
Image for post

LET’S BEGIN !!


1) Import Libraries

import os
import glob
import xml.etree.ElementTree as ET
import pandas as pd
import tensorflow as tf

2) Create customTF2, training, and data folders in your google drive

Create a folder named customTF2 in your google drive.

Create another folder named training inside the customTF2 folder ( training folder is where the checkpoints will be saved during training ).

Create another folder named data inside the customTF2 folder.


3) Create and upload your image files and its their corresponding labeled XML files.

Create a folder named images for your custom dataset images and create another folder named annotations for its corresponding PASCAL_VOC format labeled XML files.

Next, create their zip files and upload them to the customTF2 folder in your drive.

Make sure all the image files have their extension as “.jpg” only. 

Other formats like <.png> , <.jpeg> will give errors since the generate_tfrecord and xml_to_csv scripts here have only <.jpg> in them. If you have other format images, you can make changes in the scripts accordingly.

For Datasets, you can check out my Dataset Sources at the bottom of this article in the credits section. 

Collecting Images Dataset and labeling them to get their PASCAL_VOC XML annotations.

Labeling your Dataset

Input image example (Image1.jpg)

Image for post
Original Photo by Ali Pazani from Pexels


You can use any software for labeling like the labelImg tool.

Image for post
labelImg GUI for Image1.jpg

I use an open-source labeling tool called OpenLabeling with a very simple UI.

Image for post
OpenLabeling Tool GUI


Click on the link below to know more about the labeling process and other software for it:

NOTE : Garbage In = Garbage Out. Choosing and labeling images is the most important part. Try to find good quality images. The quality of the data goes a long way towards determining the quality of the result.

The output PASCAL_VOC labeled XML file looks like as shown below:

Image1.xml


4) Upload the generate_tfrecord.py file to the customTF2 folder in your drive.

You can find the generate_tfrecord.py file here


5) Mount drive and link your folder

#mount drive

from google.colab import drive
drive.mount('/content/gdrive')

# this creates a symbolic link so that now the path /content/gdrive/My Drive/ is equal to /mydrive

!ln -s /content/gdrive/My Drive/ /mydrive
!ls /mydrive

6) Clone the TensorFlow models git repository & Install TensorFlow Object Detection API

# clone the tensorflow models on the colab cloud vm

!git clone --q https://github.com/tensorflow/models.git

# navigate to /models/research folder to compile protos

%cd models/research

# Compile protos.

!protoc object_detection/protos/*.proto --python_out=.

# Install TensorFlow Object Detection API.

!cp object_detection/packages/tf2/setup.py .
!python -m pip install 

7) Test the model builder

!python object_detection/builders/model_builder_tf2_test.py

8) Navigate to /mydrive/customTF2/data/ and Unzip the images.zip and annotations.zip files into the data folder

%cd /mydrive/customTF2/data/

# unzip the datasets and their contents so that they are now in /mydrive/customTF2/data/ folder

!unzip /mydrive/customTF2/images.zip -d .
!unzip /mydrive/customTF2/annotations.zip -d .

9) Create test_labels & train_labels

Current working directory is /mydrive/customTF2/data/

Divide annotations into test_labels(20%) and train_labels(80%).


10) Create the CSV files and the “label_map.pbtxt” file

Current working directory is /mydrive/customTF2/data/

Run xml_to_csv script below to create test_labels.csv and train_labels.csv

This script also creates the label_map.pbtxt file using the classes mentioned in the xml files.

The 3 files that are created i.e. train_labels.csv, test_labels.csv, and label_map.pbtxt look like as shown below:

train_labels.csv
test_labels.csv
label_map.pbtxt

The train_labels.csv contains the name of all the train images, the classes in those images, and their annotations.

The test_labels.csv contains the name of all the test images, the classes in those images, and their annotations.

The label_map.pbtxt file contains the names of the classes from your labeled XML files. 

NOTE: I have 2 classes i.e. “with_mask” and “without_mask”.

Label map id 0 is reserved for the background label.


11) Create train.record & test.record files

Current working directory is /mydrive/customTF2/data/

Run the generate_tfrecord.py script to create train.record and test.record files

#Usage:
#!python generate_tfrecord.py output.csv output_pb.txt /path/to/images output.tfrecords

#For train.record
!python /mydrive/customTF2/generate_tfrecord.py train_labels.csv  label_map.pbtxt images/ train.record

#For test.record
!python /mydrive/customTF2/generate_tfrecord.py test_labels.csv  label_map.pbtxt images/ test.record

The total number of image files is 1370. Since we divided the labels into two categories viz. train_labels(80%) and test_labels(20%), the number of files for “train.record” is 1096, and the number of files for “test.record” is 274.


12) Download pre-trained model checkpoint

Current working directory is /mydrive/customTF2/data/

You can choose any model for training depending upon your data and requirement. Read this blog for more info on this. The official list of detection model checkpoints for TensorFlow 2.x can be found here.

In this tutorial, I will use the ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8 model.

# Download the pre-trained model ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8.tar.gz into the data folder & unzip it

!wget http://download.tensorflow.org/models/object_detection/tf2/20200711/ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8.tar.gz

!tar -xzvf ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8.tar.gz

13) Get the model pipeline config file, make changes to it and put it inside the data folder

Current working directory is /mydrive/customTF2/data/

Download ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8.config from /content/models/research/object_detection/configs/tf2. Make the required changes to it and upload it to the /mydrive/custom/data folder.

OR

Edit the config file from /content/models/research/object_detection/configs/tf2 in colab vm and copy the edited config file to the /mydrive/customTF2/data folder.

You need to make the following changes:

  • change num_classes to the number of your classes.
  • change test.record path, train.record path & labelmap path to the paths where you have created these files (paths should be relative to your current working directory while training).
  • change fine_tune_checkpoint to the path of the directory where the downloaded checkpoint from step 12 is.
  • change fine_tune_checkpoint_type with value classification or detection depending on the type.
  • change batch_size to any multiple of 8 depending upon the capability of your GPU. (eg:- 24,128,…,512). Mine is set to 64.
  • change num_steps to the number of steps you want the detector to train.
# copy the edited config file from the configs/tf2 directory to the data/ folder in your drive

!cp /content/models/research/object_detection/configs/tf2/ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8.config /mydrive/customTF2/data

The workspace at this point:

There are many data augmentation options that you can add. Check the full list here. For beginners, the above changes are sufficient.

Data Augmentation Suggestions

First, you should train the model using the sample config file with the above basic changes and see how well it does. If you are overfitting, then you might want to do some more image augmentations.

  • In the sample config file: random_horizontal_flip & ssd_random_crop are added by default. You could try adding these as well:

from train_config {}:

data_augmentation_options {
    random_adjust_contrast {
    }
  }
  data_augmentation_options {
    random_rgb_to_gray {
    }
  }
  data_augmentation_options {
    random_vertical_flip {
    }
  }
  data_augmentation_options {
    random_rotation90 {
    }
  }
  data_augmentation_options {
    random_patch_gaussian {
    }
  }

Note: Each image augmentation will increase the training time drastically.

In model {} > ssd {} > box_predictor {}: set use_dropout to true This will be helpful to counter overfitting.

In eval_config : {} set the number of testing images you have in num_examples and remove max_eval to evaluate indefinitely

eval_config: {
  num_examples: 274 # set this to the number of test images we divided earlier
  num_visualizations: 20 # the number of visualization to see in tensorboard
}

14) Load Tensorboard

%load_ext tensorboard
%tensorboard --logdir '/content/gdrive/MyDrive/customTF2/training'

15) Train the model

Navigate to the object_detection folder in Colab VM

%cd /content/models/research/object_detection

15 (a) Training using model_main_tf2.py

Here {PIPELINE_CONFIG_PATH} points to the pipeline config and {MODEL_DIR} points to the directory in which training checkpoints and events will be written.

For best results, you should stop the training when the loss is less than 0.1 if possible, else train the model until the loss does not show any significant change for a while. (Try to get the loss as low as possible. The ideal loss should be below 0.05).

# Run the command below from the content/models/research/object_detection directory

#PIPELINE_CONFIG_PATH={path to pipeline config file}
#MODEL_DIR={path to model directory}
#!python model_main_tf2.py --pipeline_config_path=${PIPELINE_CONFIG_PATH} --model_dir=${MODEL_DIR} --alsologtostderr

!python model_main_tf2.py --pipeline_config_path=/mydrive/customTF2/data/ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8.config --model_dir=/mydrive/customTF2/training --alsologtostderr

IMPORTANT: The output will normally look like it has “frozen”, but DO NOT rush to cancel the process. The training outputs logs only every 100 steps by default, therefore if you wait for a while, you should see a log for the loss at step 100.

The time you should wait can vary greatly, depending on whether you are using a GPU and the chosen value for batch_size in the config file, so be patient.

15 (b) Evaluation using model_main_tf2.py

Here {CHECKPOINT_DIR} points to the directory with checkpoints produced by the training job. Evaluation events are written to {MODEL_DIR/eval}.

For best results, you should stop the training when the loss is less than 0.1 if possible, else train the model until the loss does not show any significant change for a while. (Try to get the loss as low as possible. The ideal loss should be below 0.05).

# Run the command below from the content/models/research/object_detection directory

#!python model_main_tf2.py --pipeline_config_path=${PIPELINE_CONFIG_PATH} --model_dir=${MODEL_DIR} --checkpoint_dir=${CHECKPOINT_DIR} --alsologtostderr

#PIPELINE_CONFIG_PATH={path to pipeline config file}

#MODEL_DIR={path to model directory}

#CHECKPOINT_DIR=${MODEL_DIR}


!python model_main_tf2.py --pipeline_config_path=/mydrive/customTF2/data/ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8.config --model_dir=/mydrive/customTF2/training/ --checkpoint_dir=/mydrive/customTF2/training/ --alsologtostderr

RETRAINING THE MODEL ( in case you get disconnected )

If you get disconnected or lose your session on Colab VM, you can start your training where you left off as the checkpoint is saved on your drive inside the training folder. To restart the training simply run steps 1, 5, 6, 7, 14 and 15.

Note that since we have all the files required for training like the record files, our edited pipeline config file, the labelmap file, and the model checkpoint folder, therefore we do not need to create these again.

The model_main_tf2.py script saves the checkpoint every 1000 steps. The training automatically restarts from the last saved checkpoint itself.

However, if you see that it doesn’t restart training from the last checkpoint you can make 1 change in the pipeline config file. Change fine_tune_checkpoint to where your latest trained checkpoints have been written and have it point to the latest checkpoint as shown below:

fine_tune_checkpoint: "/mydrive/customTF2/training/ckpt-X" (where ckpt-X is the latest checkpoint)

Read this TensorFlow Object Detection API tutorial to know more about the training process for TF2.


16) Test your trained custom object detection model

Export inference graph

Current working directory is /content/models/research/object_detection

!python exporter_main_v2.py --trained_checkpoint_dir=/mydrive/customTF2/training --pipeline_config_path=/content/gdrive/MyDrive/customTF2/data/ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8.config --output_directory /mydrive/customTF2/data/inference_graph

Test your trained custom object detection model on images

Current working directory is /content/models/research/object_detection

This step is optional.
# Different font-type and font-size for labels text

!wget https://freefontsdownload.net/download/160187/arial.zip
!unzip arial.zip -d .

%cd utils/
!sed -i "s/font = ImageFont.truetype('arial.ttf', 24)/font = ImageFont.truetype('arial.ttf', 50)/" visualization_utils.py
%cd ..

Test your trained model


NOTE:

The dataset I have collected for mask detection contains mostly close-up images. For more long-shot images you can search online. There are many sites where you can download labeled and unlabeled datasets. I have given a few links at the bottom under Dataset Sources. I have also given a few links for mask datasets. Some of them have more than 10,000 images.

Though there are certain tweaks and changes we can make to our training config file or add more images to the dataset for every type of object class through augmentation, we have to be careful so that it does not cause overfitting which affects the accuracy of the model.

For beginners, you can start simply by using the config file I have uploaded on my GitHubI have also uploaded my mask images dataset along with the PASCAL_VOC format labeled text files, which although might not be the best but will give you a good start on how to train your own custom object detector using an SSD model. You can find a labeled dataset of better quality or an unlabeled dataset and label it yourself later.

Image for post
Original Video by Max Fischer from Pexels

My GitHub

Files for training

Train Custom Object Detection model using TF 2

My mask dataset

https://www.kaggle.com/techzizou/labeled-mask-dataset-pascal-voc-format

My Colab notebook for this

Custom Object Detection TensorFlow 2.x

Check out my Youtube Video on this

Coming Soon!


CREDITS

Documentation / References

Dataset Sources

You can download datasets for many objects from the sites mentioned below. These sites also contain images of many classes of objects along with their annotations/labels in multiple formats such as the YOLO_DARKNET txt files and the PASCAL_VOC xml files.

Mask Dataset Sources

More Mask Datasets


Leave a Reply