TensorFlow 2 Object Detection API using Custom Dataset
By: Renu Khandelwal
Tags: Tensorboard, Python, Object Detection, Tensorflow2, Deep Learning

TensorFlow 2 Object Detection API using Custom Dataset

Everything you should know about using TF2 Object Detection API on a Custom Dataset, including common issues faced and their solutions

Here you will go step by step to perform object detection on a custom dataset using TF2 Object Detection API and some of the issues and resolutions.

Object Detection using TF2 Object Detection API on Kangaroo dataset

The custom dataset is available here.

TensorFlow 2 Object detection model is a collection of detection models pre-trained on the COCO 2017 dataset.

Tensorflow 2 Object Detection API in this article will identify all the kangaroo’s present in an image or video, along with their locations. Locations of kangaroo will be depicted by drawing bounding boxes.

Steps for performing Object Detection on a Custom Dataset using TF2 Object Detection API

Step1 : Setup for TensorFlow2 Object Detection API

Clone the TensorFlow Models repository

git clone https://github.com/tensorflow/models.git

You should have the following directory structure.

TF2 Object Detection API folder structure

You now have the code for TF2 Object detection API; install the API using the protobuf compiler. Ensure you are using version 3 or higher of the protobuf compiler.

Add the following paths to the PYTHONPATH environment variable for the TF2 Object detection API to work.

Path to …\models\research Path to …\models\research\object_detection

Some common issues and resolutions for the protobuf compiler can be found here.

Once the environment variable PYTHONPATH is set up, you can execute the following command to convert all the .proto files in the \models\research\object_detection\protos folder .py files.

%%bash

cd models/research/

protoc object_detection\protos\*.proto — python_out=.

cp object_detection\packages\tf2\setup.py .

python -m pip install .

Ideally, pycocotools gets installed when installing the TF2 Object Detection API, but if you have issues and it is not installed, then clone the cocoapi repository as shown below.

Step 2: Set up the directory structure for the Custom dataset

Directory setup for a custom dataset

Explanation of the contents of the directories required for object detection for training on a custom dataset.

annotations: will store the TFRecord files for train and test dataset along with the label_map.pbtxt. You can also store the .csv files used for generating the TFRecord files. TFRecord contains a sequence of records, including annotation details. TFRecord files are read sequentially.

images: will store the image file and annotations XML files for both the train and the test dataset.

models: will contain a sub-directory for each of the training jobs along with the pipeline.config file. If you are using different object detection models like SSD and EffeicientDet, you can create two sub-folders for SSD and one for EfficientDet, and each sub-folder will have its own pipeline.config file.

pretrained_models: will contain the downloaded pre-trained models and will be used as a starting point for our custom dataset training.

exported_model: will store the exported versions of the trained model on the custom dataset.

Step 3: Create a Label Map file

TF2 object detection API requires a label_map.pbtxt file containing the mapping of the objects you want to detect mapped to an integer value, as shown below. The Object Detection API uses label_map.txt file during training as well as during inference.

item {

id: 1

name: 'kangaroo'

}

Place this file in the annotations folder, as shown below.

Step 4: Prepare the Custom Dataset

In this article, we are using a kangaroo dataset that is already annotated. If you have images, then you can use tools like labelImg for annotating the images.

We will read the annotation files and create a data frame that will be used for generating the TFRecord.

The images and the annotated XML files should be placed under the images folder within the respective train or test directories, as shown below.

Images and XML files containing annotations

One of the issues I faced was having a corrupted image file, and I was getting an error “Invalid JPEG data or crop window”. I used the following code to find the corrupted image file.

You can use any one of the libraries io or cv2

## To check if the Jpeg are good

from skimage import io

import cv2

files=[]

files=glob.glob(TRAIN_IMAGE_FILE+ '\*.jpg')

for i in range(len(files)):



try:

_ = io.imread(files[i])

img = cv2.imread(files[i])



except Exception as e:

print(e)

print(files[i])

Creating TFRecord based on the annotated XML file

Create the TFRecord for the training and test annotated XML files using the following code

# importing required libraraies

import os

import glob

import pandas as pd

import io

import xml.etree.ElementTree as ET

import argparse

import tensorflow.compat.v1 as tf

from PIL import Image

from object_detection.utils import dataset_util, label_map_util

from collections import namedtuple # Set the folder name for the source annotated XML files and folder #to store the TFRecord Record file

LABEL_MAP_FILE=r'\Custom_OD\Workspace\Annotations\label_map.pbtxt'

TRAIN_XML_FILE=r'\Custom_OD\Workspace\images\Train'

TRAIN_TF_RECORD_DIR=r'\Custom_OD\Workspace\Annotations\train.record'

TEST_XML_FILE=r'\Custom_OD\Workspace\images\test'

TEST_TF_RECORD_DIR=r'\Custom_OD\Workspace\Annotations\test.record' #Create a dictionary for the labels or objects

label_map = label_map_util.load_labelmap(LABEL_MAP_FILE)

label_map_dict = label_map_util.get_label_map_dict(label_map) #convert the object annotation from XML file to a dataframe

def xml_to_df(path):

"""Iterates through all .xml files conatining Annotation in a given directory and combines

them in a single Pandas dataframe.

Parameters:

----------

path : str

The path containing the .xml files

Returns

-------

Pandas DataFrame

The produced dataframe

""" xml_list = []

for xml_file in glob.glob(path + '/*.xml'):

tree = ET.parse(xml_file)

root = tree.getroot()

for member in root.findall('object'):

value = (root.find('filename').text,

int(root.find('size')[0].text),

int(root.find('size')[1].text),

member[0].text,

int(member[4][0].text),

int(member[4][1].text),

int(member[4][2].text),

int(member[4][3].text)

)

xml_list.append(value)

column_name = ['filename', 'width', 'height',

'class', 'xmin', 'ymin', 'xmax', 'ymax']

xml_df = pd.DataFrame(xml_list, columns=column_name)

return xml_df #pass the label and get its equivalent integer

def class_label_to_int(row_label):

return label_map_dict[row_label] #Split filename and the annotations details for all the xml files

def split(df, group):

data = namedtuple('data', ['filename', 'object'])

gb = df.groupby(group)

return [data(filename, gb.get_group(x)) for filename, x in zip(gb.groups.keys(), gb.groups)] # Create the TF Record files

def create_tf_example(group, path):

with tf.io.gfile.GFile(os.path.join(path, '{}'.format(group.filename)), 'rb') as fid:

encoded_jpg = fid.read()



encoded_jpg_io = io.BytesIO(encoded_jpg)

image = Image.open(encoded_jpg_io)

width, height = image.size

filename = group.filename.encode('utf8')

image_format = b'jpg'

xmins = []

xmaxs = []

ymins = []

ymaxs = []

classes_text = []

classes = [] for index, row in group.object.iterrows():

xmins.append(row['xmin'] / width)

xmaxs.append(row['xmax'] / width)

ymins.append(row['ymin'] / height)

ymaxs.append(row['ymax'] / height)

classes_text.append(row['class'].encode('utf8'))

classes.append(class_label_to_int(row['class'])) tf_example = tf.train.Example(features=tf.train.Features(feature={

'image/height': dataset_util.int64_feature(height),

'image/width': dataset_util.int64_feature(width),

'image/filename': dataset_util.bytes_feature(filename),

'image/source_id': dataset_util.bytes_feature(filename),

'image/encoded': dataset_util.bytes_feature(encoded_jpg),

'image/format': dataset_util.bytes_feature(image_format),

'image/object/bbox/xmin': dataset_util.float_list_feature(xmins),

'image/object/bbox/xmax': dataset_util.float_list_feature(xmaxs),

'image/object/bbox/ymin': dataset_util.float_list_feature(ymins),

'image/object/bbox/ymax': dataset_util.float_list_feature(ymaxs),

'image/object/class/text': dataset_util.bytes_list_feature(classes_text),

'image/object/class/label': dataset_util.int64_list_feature(classes),

}))

return tf_example #Generating the train TFRecord

writer = tf.python_io.TFRecordWriter(TRAIN_TF_RECORD_DIR)

for group in grouped_train:

tf_example = create_tf_example(group, TRAIN_XML_FILE)

writer.write(tf_example.SerializeToString())

writer.close()

print('Successfully created the TFRecord file: {}'.format(TRAIN_TF_RECORD_DIR)) # Generating the test TFRecord

writer = tf.python_io.TFRecordWriter(TEST_TF_RECORD_DIR)

for group in grouped_test:

tf_example = create_tf_example(group, TEST_XML_FILE)

writer.write(tf_example.SerializeToString())

writer.close()

print('Successfully created the TFRecord file: {}'.format(TEST_TF_RECORD_DIR))

Now the folder structure with files should be as shown below.

Step 5: Download the Pretrained Models

Download your choice of pre-trained object detection model from here. You will have a tar file like ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8.tar.gz. Extract the contents of the file in the pretrained_models folder.

The extracted folder should have a checkpoint folder, saved_model folder, and pipeline.config file for the model.

You can download multiple pre-trained models in the pretrained_models folder, as shown below.

Step 6: Configuring the Training Pipeline by updating the pipeline.config

Copy the pre-trained model’s pipeline.config file to the model’s directory and edit a few settings to prepare the object detection API for the custom dataset.

num_classes: update num_classes to match the classes/objects in your dataset and mentioned in label_map.pbtxt

model {

ssd {

num_classes: 1

image_resizer {

fixed_shape_resizer {

height: 320

width: 320

}

}

}

batch_size: update batch_size based on the system's memory you will be training on. The higher batch_size will need a higher memory.

train_config {

batch_size: 6

data_augmentation_options {

random_horizontal_flip {

}

}

Guidance for batch_size as mentioned here

Max batch size= available GPU memory(in bytes) / 4 / (size of tensors + trainable parameters)

fine_tune_checkpoint : mention the path to the checkpoint of the pre-trained model

: mention the path to the checkpoint of the pre-trained model num_steps: number of training steps you want to train the object detection API on

number of training steps you want to train the object detection API on fine_tune_checkpoint_type : specify “ detection ” when performing object detection and specify “ classification ” when performing image classification

: specify “ ” when performing object detection and specify “ ” when performing image classification use_bfloat16: set this to false if you are not training on a TPU

fine_tune_checkpoint: "\\Custom_OD\\Workspace\\pretrained_model\\SSD_640\\ckpt-0"

num_steps: 300000

startup_delay_steps: 0.0

replicas_to_aggregate: 8

max_number_of_boxes: 100

unpad_groundtruth_tensors: false

fine_tune_checkpoint_type: "detection"

use_bfloat16: true

fine_tune_checkpoint_version: V2

Training parameters: update the label_map_path to point to the label_map.pbtxt file in the annotations folders and update the input_path to point to the training TFRecord file in the annotations folder

train_input_reader {

label_map_path: "\\Custom_OD\\Workspace\\Annotations\\label_map.pbtxt"

tf_record_input_reader {

input_path: "\\Custom_OD\\Workspace\\Annotations\\train.record"

}

}

Evaluation parameters: update the label_map_path to point to the label_map.pbtxt file in the annotations folders and update the input_path to point to the test TFRecord file in the annotations folder

eval_input_reader {

label_map_path: "\\Custom_OD\\Workspace\\Annotations\\label_map.pbtxt"

shuffle: false

num_epochs: 1

tf_record_input_reader {

input_path: "\\Custom_OD\\Workspace\\Annotations\\test.record"

}

}

Another issue I encountered was the error “Invoked with: None, 524288”. The resolution is to check all of the above pipeline.config file parameters and ensure a valid value for each of the file parameters.

Step 7: Training the Object detection model on the custom dataset

To start training object detection API using any of the pre-trained models on a custom dataset, you need to copy the model_main_tf2.py file from the object detection API to the workspace folder as shown below.

Execute the following command on your command prompt to start the training

python model_main_tf2.py

--model_dir=models\SSD_640

--pipeline_config_path=models\SSD_640\pipeline.config

model_dir : is the path where the training process will store the checkpoint files.

: is the path where the training process will store the checkpoint files. pipeline_config_path is the path where the pipeline.config file is stored

The model will be trained based on the num_steps specified in the pipeline.config file and the time to train will depend on your host platform/machine's computation capability.

As the model is trained, it will occasionally create checkpoint files and place them in the model_dir specified for training the model.

Evaluating the Object Detection API for the custom dataset

Update the pipeline.config file to point to the right checkpoint as shown below

fine_tune_checkpoint: "\\Custom_OD\\Workspace\\models\\SSD_640\\ckpt-6"

To evaluate the object detection trained model execute the following command from the command prompt

python model_main_tf2.py

--model_dir=models\SSD_640

--pipeline_config_path=models\SSD_640\pipeline.config --checkpoint_dir=models\SSD_640

model_dir : the path will contain the events files generated as part of the evaluation

: the path will contain the events files generated as part of the evaluation pipeline_config_path : is the path where the pipeline.config file is stored

: is the path where the pipeline.config file is stored checkpoint_dir: is the path where the training process has stored the checkpoint files.

After the evaluation is complete, you can see the events files in the path specified in the model_dir.

When the evaluation process runs, it periodically checks for a new checkpoint file in the checkpoint_dir. Once the evaluation process is complete, then you can see the following output.

Evaluation

Monitoring the Object Detection Training

The Object Detection API training on a custom dataset can be monitored using TensorBoard. The events file generated as part of the evaluation will display the evaluation metrics on TensorBoard.

To start the TensorBoard, execute the following command on the command prompt

tensorboard --logdir=models\SSD_640

log_dir: specify the path that contains the checkpoints and the evaluation events

once the command is executed you will see the URL as shown below

Exporting the Trained Model

You will now need to export the trained model to be used for Object detection on your custom dataset.

First, copy the exporter_main_v2.py file from the object detection API code in the models/research/object_detection folder to your workspace, as shown below.

To export a trained object detection model, execute the following command from the command prompt

python exporter_main_v2.py

--input_type=image_tensor

--trained_checkpoint_dir=models\SSD_640

--pipeline_config_path=models\SSD_640\pipeline.config

--output_directory=exported_model

input_type : is the input node type of the inference graph. options are image_tensor, encoded_image_string_tensor, and tf_example. image_tensor is a uint8 4-D tensor.

: is the input node type of the inference graph. options are image_tensor, encoded_image_string_tensor, and tf_example. image_tensor is a uint8 4-D tensor. trained_checkpoint : the path to the folder containing the trained checkpoint

: the path to the folder containing the trained checkpoint pipeline_config_path : path for the pipeline config file

: path for the pipeline config file output_directory: the exported model will be available in the path specified in the output_directory

Making Inferences Using the Trained Object Detection Model

you now have a trained model that you have exported to a folder that can make inferences.

#Import the required libraries for Object detection infernece

import time

import tensorflow as tf

from object_detection.utils import label_map_util

from object_detection.utils import visualization_utils as viz_utils

import os

import cv2

import matplotlib.pyplot as plt

%matplotlib inline # setting min confidence threshold

MIN_CONF_THRESH=.6 #Loading the exported model from saved_model directory

PATH_TO_SAVED_MODEL =r'Custom_OD\Workspace\exported_model\saved_model'

print('Loading model...', end='')

start_time = time.time() # LOAD SAVED MODEL AND BUILD DETECTION FUNCTION

detect_fn = tf.saved_model.load(PATH_TO_SAVED_MODEL)

end_time = time.time()

elapsed_time = end_time - start_time

print('Done! Took {} seconds'.format(elapsed_time)) # LOAD LABEL MAP DATA

PATH_TO_LABELS=r'\Custom_OD\Workspace\Annotations\label_map.pbtxt'

category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True) #Image file for inference

IMAGE_PATH=r'\Custom_OD\Workspace\images\test\00178.jpg' def load_image_into_numpy_array(path):

"""Load an image from file into a numpy array.

Puts image into numpy array of shape (height, width, channels), where channels=3 for RGB to feed into tensorflow graph.

Args:

path: the file path to the image

Returns:

uint8 numpy array with shape (img_height, img_width, 3)

"""

return np.array(cv2.cvtColor(cv2.imread(path), cv2.COLOR_BGR2RGB))

image_np = load_image_into_numpy_array(IMAGE_PATH) # Running the infernce on the image specified in the image path

# The input needs to be a tensor, convert it using `tf.convert_to_tensor`.

input_tensor = tf.convert_to_tensor(image_np)

# The model expects a batch of images, so add an axis with `tf.newaxis`.

input_tensor = input_tensor[tf.newaxis, ...]

detections = detect_fn(input_tensor)



# All outputs are batches tensors.

# Convert to numpy arrays, and take index [0] to remove the batch dimension.

# We're only interested in the first num_detections.

num_detections = int(detections.pop('num_detections'))

detections = {key: value[0, :num_detections].numpy()

for key, value in detections.items()}

detections['num_detections'] = num_detections # detection_classes should be ints.

detections['detection_classes'] = detections['detection_classes'].astype(np.int64) #print(detections['detection_classes'])

image_np_with_detections = image_np.copy() viz_utils.visualize_boxes_and_labels_on_image_array(

image_np_with_detections,

detections['detection_boxes'],

detections['detection_classes'],

detections['detection_scores'],

category_index,

use_normalized_coordinates=True,

max_boxes_to_draw=200,

min_score_thresh=MIN_CONF_THRESH,

agnostic_mode=False) plt.figure()

plt.imshow(image_np_with_detections)

print('Done')

plt.show()

Conclusion:

Use TensorFlow2 object detection API for detecting objects in your custom dataset step by step.

Download the object detection API and coco API,

Set up the folder structure as shown in the article

For the Custom dataset, create a customized label_map.pbtxt file in the annotations folder containing all the objects you wish to detect in your custom dataset and assign an integer for all the classes.

Prepare the custom dataset for training and testing and convert the XML annotations to TFRecord files for train and test and store it in annotations folders

Download the pre-trained model fo your choice from the model zoo

Update the pipeline.config file for mapping to a custom dataset

Once the model is trained you can monitor the training and evaluation metrics using TensorBoard

Export the trained model that will be used for making inferences on images or videos

References:

https://github.com/tensorflow/models