A simple example to how to use Tensorflow detection API on a custom dataset

_{^{Detector in action (from https://youtu.be/eYUohOTqn7g)}}

Download and install labelImg tool

• Put your images in dataset/train/images folder
• Create a dataset/train/annotations/ folder
• Open Dir in labelImg tool to dataset/train/images folder
• Change Save Dir in labelImg tool to dataset/train/annotations/ folder
• Start boring annotation task 🤷‍♀️

For the laziest, you can also use an existing dataset. Use the script convert_dataset_format.py to convert from a YOLO dataset to a PASCAL VOC style dataset.

Download and install the tensorflow object detection preferably in a conda env

conda create --name tf_detection python=3.7 -y
conda activate tf_detection

Do not forget to add it to the PYTHONPATH

Set the TF_MODELS_PATH env var to the Tensorflow detection folder. If you cloned in your home directory export TF_MODELS_PATH=~/models or cd into the directory and do export TF_MODELS_PATH=$(pwd)

Split test train

python split_train_test.py

Create the TFRecords

python create_pascal_tf_record.py --output_path=./train.record --data_dir=dataset/train
python create_pascal_tf_record.py --output_path=./test.record --data_dir=dataset/test

The model used in this example is the ssdlite_mobilenet_v2_coco_2018_05_09, you can change it with whatever you like

python download_model.py

If you changed the model, download also the adequate pipeline config here and edit the file by replacing all the PATH_TO_BE_CONFIGURED.

Run this command to begin the training

python ${TF_MODELS_PATH}/research/object_detection/model_main.py \
   --pipeline_config_path=ssdlite_mobilenet_v2_coco_modified.config \
   --model_dir=./finetuned_model \
   --alsologtostderr \
   --num_train_steps=30000 \
   --num_eval_steps=100

You can follow the training using Tensorboard

tensorboard --logdir=finetuned_model

After the training completion, export the graph for inference

INPUT_TYPE=image_tensor
PIPELINE_CONFIG_PATH=ssdlite_mobilenet_v2_coco_modified.config
TRAINED_CKPT_PREFIX=finetuned_model/model.ckpt-30000
EXPORT_DIR=finetuned_model/export_frozen
python ${TF_MODELS_PATH}/research/object_detection/export_inference_graph.py \
    --input_type=${INPUT_TYPE} \
    --pipeline_config_path=${PIPELINE_CONFIG_PATH} \
    --trained_checkpoint_prefix=${TRAINED_CKPT_PREFIX} \
    --output_directory=${EXPORT_DIR}

To run a inference in a image

python inference.py