The easiest way to Train a Custom Object Detection Model Using TensorFlow Object Detection API
From: https://makeoptim.com/en/deep-learning/yiai-object-detection
Foreword
This article will introduce the concept of Object Detection, and explain how to use TensorFlow Object Detection API to train a custom object detector through cases, including data set collection and processing, TensorFlow Object Detection API installation, and model training.
The case effect is shown in the figure below:
Object Detection
As shown in the figure above, the problem that Image Classification solves is what is the object in the picture, and Object Detection can identify what objects and the location(coordinate) of the object in the picture.
Location
The location information of Object Detection generally has two formats:
- Polar coordinates (xmin, ymin, xmax, ymax): xmin,ymin: the minimum value of x,y coordinates; xmin,ymin: the maximum value of x,y coordinates
- Center point: (x_center, y_center, w, h): x_center, y_center: the coordinates of the center point of the target detection frame; w,h: the width and height of the target detection frame
Note: The upper left corner of the picture is the origin (0,0)
Milestone
Traditional method (region proposal + manual feature extraction + classifier)
HOG+SVM、DPM
Region Proposal+CNN(Two-stage)
R-CNN, SPP-NET, Fast R-CNN, Faster R-CNN
End-to-End(One-stage)
YOLO、SSD
TensorFlow Object Detection API
The TensorFlow Object Detection API is an open-source framework built on top of TensorFlow that makes it easy to construct, train and deploy object detection models. In addition, the TensorFlow Object Detection API also provides Model Zoo to facilitate our selection and switching of pre-trained models.
Install dependencies
Use the following command to check if the installation is successful.
$ conda --version
conda 4.9.2
$ protoc --version
libprotoc 3.17.1
Install API
The official installation steps provided by the TensorFlow Object Detection API are more cumbersome. The author wrote a script to install directly in one step.
Execute git clone https://github.com/CatchZeng/object-detection-api.git
to download the repo, then go to the directory of the repo (hereafter referred to as the oda
repo) and execute the following command, if you see the following output, indicating that the installation was successful.
$ conda create -n od python=3.8.5 && conda activate od && make install......
----------------------------------------------------------------------
Ran 24 tests in 21.869sOK (skipped=1)
Note: If you don’t want to use conda, you can install it directly with `make install` in your python environment, for example, use it in colab.
Note: The update of
cudaDNN
andtoolkit
may not be as fast as TensorFlow. Therefore, if your machine has a GPU, after the installation is complete, you need to downgrade TensorFlow to the version supported bycudaDNN
andtoolkit
in order to support GPU training. Take2.4.1
as an example:
$ pip install --upgrade tf-models-official==2.4.0
$ pip install --upgrade tensorflow==2.4.1
Note: If the installation fails, you can refer to the detailed steps in the official document.
Create workspace
Note: !!! From here, please make sure to execute under the environment of
conda od
.
$ conda activate od
$ conda env list
# conda environments:
#
od * /Users/catchzeng/.conda/envs/od
tensorflow /Users/catchzeng/.conda/envs/tensorflow
base /Users/catchzeng/miniconda3
Go to the oda
repo directory and execute the following command to create the workspace directory structure.
Note:
SAVE_DIR
is the directory to save the workspace, andNAME
is the name of the workspace.
$ make workspace-box SAVE_DIR=workspace NAME=test└── workspace
└── test
├── Makefile
├── annotations: Store the labeled data set, including val.record, train.record, label_map.pbtxt
├── convert_quant_lite.py: the script to quantize tflite model
├── export_tflite_graph_tf2.py: the script to export tflite model
├── exported-models: Store the exported model after training
├── exporter_main_v2.py: the script to export model
├── images: Data set images and xml annotations
│ ├── test: Store manual verification images
│ ├── train: Training set images and xml annotations
│ └── val: Evaluation set images and xml annotations
├── model_main_tf2.py:the script to train model
├── models: Custom model
├── pre-trained-models: pre-trained models
└── test_images.py: the script to verify the image manually
Dataset
Images
I like to drink tea. Today I will use a cup, teapot, humidifier as examples.
Put the collected pictures into the three subdirectories of images
in the project directory.
Note: This case is just to show how to train the Object Detection model, so the dataset is relatively small. In the production, remember to collect as many datasets as possible.
Annotation
After collecting the pictures, you need to annotate the images in the training and evaluation sets.
We choose LabelImg as the annotation tool.
Install LabelImg according to the instructions of installation, then execute labelImg
to select the train
and val
folders for annotation.
After the annotation is completed, the xml
annotation file corresponding to the image will be generated, as shown below:
workspace/test/images
├── test
│ ├── 15.jpg
│ └── 16.jpg
├── train
│ ├── 1.jpg
│ ├── 1.xml
│ ├── 10.jpg
│ ├── 10.xml
│ ├── 2.jpg
│ ├── 2.xml
│ ├── 3.jpg
│ ├── 3.xml
│ ├── 4.jpg
│ ├── 4.xml
│ ├── 5.jpg
│ ├── 5.xml
│ ├── 6.jpg
│ ├── 6.xml
│ ├── 7.jpg
│ ├── 7.xml
│ ├── 8.jpg
│ ├── 8.xml
│ ├── 9.jpg
│ └── 9.xml
└── val
├── 11.jpg
├── 11.xml
├── 12.jpg
├── 12.xml
├── 13.jpg
├── 13.xml
├── 14.jpg
└── 14.xml
Create TFRecord
TensorFlow Object Detection API only supports TFRecord format, therefore, the dataset needs to be converted.
Go to the workspace directory (cd workspace/test
), and then execute make gen-tfrecord
, it will generate label_map.pbtxt
and TFRecord format dataset in the annotations
folder.
$ make gen-tfrecord
python gen_label_map.py
unsorted: ['cup', 'teapot', 'humidifier']
sorted: ['cup', 'humidifier', 'teapot']
item {
id: 1
name: 'cup'
}
item {
id: 2
name: 'humidifier'
}
item {
id: 3
name: 'teapot'
}
python generate_tfrecord.py \
-x images/train \
-l annotations/label_map.pbtxt \
-o annotations/train.record
Successfully created the TFRecord file: annotations/train.record
python generate_tfrecord.py \
-x images/val \
-l annotations/label_map.pbtxt \
-o annotations/val.record
Successfully created the TFRecord file: annotations/val.recordannotations
├── label_map.pbtxt
├── train.record
└── val.record# label_map.pbtxt
item {
id: 1
name: 'cup'
}
item {
id: 2
name: 'humidifier'
}
item {
id: 3
name: 'teapot'
}
Model training
Note: !!! From here, please make sure you have go to the workspace directory (
cd workspace/test
).
Download the pre-trained model
Select the appropriate model from Model Zoo, download and unzip it, and put it in workspace/test/pre-trained-models
.
If you choose SSD MobileNet V2 FPNLite 320x320
, you can execute the following command to automatically download and decompress
$ make dl-model
The directory structure is as follows:
└── test
└── pre-trained-models
└── ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8
├── checkpoint
├── pipeline.config
└── saved_model
Configure training pipeline
Create a corresponding model folder in the models
directory, for example: ssd_mobilenet_v2_fpnlite_320x320
, and copy pre-trained-models/ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8/pipeline.config
.
└── test
├── models
│ └── ssd_mobilenet_v2_fpnlite_320x320
│ └── pipeline.config
└── pre-trained-models
Among them, pipeline.config
needs to be modified according to the project as follows
model {
ssd {
num_classes: 3 # Modify to the number of objects that need to be identified.
......
}
train_config {
batch_size: 8 # Here you need to adjust the size according to your own computer performance
......
optimizer {
momentum_optimizer {
learning_rate {
cosine_decay_learning_rate {
learning_rate_base: 0.07999999821186066
total_steps: 10000 # Modify to the total number of steps you want to train
warmup_learning_rate: 0.026666000485420227
warmup_steps: 1000
}
}
momentum_optimizer_value: 0.8999999761581421
}
use_moving_average: false
}
fine_tune_checkpoint: "pre-trained-models/ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8/checkpoint/ckpt-0" # Modify the path to the pre-trained model
num_steps: 10000 # Modify to the total number of steps you want to train
startup_delay_steps: 0.0
replicas_to_aggregate: 8
max_number_of_boxes: 100
unpad_groundtruth_tensors: false
fine_tune_checkpoint_type: "detection" # Here needs to be modified to detection, because we are doing object detection
fine_tune_checkpoint_version: V2
}
train_input_reader {
label_map_path: "annotations/label_map.pbtxt" # Modify to the annotations path
tf_record_input_reader {
input_path: "annotations/train.record" # Modify the path to the training set
}
}
eval_config {
metrics_set: "coco_detection_metrics"
use_moving_averages: false
}
eval_input_reader {
label_map_path: "annotations/label_map.pbtxt" # Modify to the annotations path
shuffle: false
num_epochs: 1
tf_record_input_reader {
input_path: "annotations/val.record" # Modify the path to the evaluation set
}
}
Training model
$ make train
Note: If the following problems occur
ValueError: numpy.ndarray size changed, may indicate binary incompatibility. Expected 88 from C header, got 80 from PyObject
You can execute the following command to solve it.
pip uninstall numpy
pip install numpy
Model export and conversion
- Saved Model
$ make export
- TFLite Model
$ make export-lite
- Convert TFLite model
$ make convert-lite
- Quantize TFLite Model
$ make convert-quant-lite
Note: For the above command, you can add
-640
to indicate the use of the model ofSSD MobileNet V2 FPNLite 640x640
, for example:make train
->make train-640
Testing
After executing make export
to export the model, and put the test image in the images/test
folder, and then execute python test_images.py
to output the annotated image to images/test_annotated
.
Summary
This article has gone through the entire process of Object Detection through cases, hoping to help you quickly master the ability to train custom Object Detectors.
The code and dataset of the case have been placed in https://github.com/CatchZeng/object-detection-api.
The following articles will bring you the principles of Object Detection, Popular Object detection networks, and Image Segmentation. That’s it for this article, see you next.
Further reading
- The easiest way to Train a Custom Image Segmentation Model Using TensorFlow Object Detection API Mask R-CNN
- How to deploy an image segmentation model service