Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
TensorFlow Object Detection Model Training

TensorFlow Object Detection Model Training

This is a summary of this nice tutorial.

Prerequisites

  1. Install TensorFlow.
  2. Download the TensorFlow models repository.

Annotating images and serializing the dataset

All the scripts mentioned in this section receive arguments from the command line and have help messages through the -h/--help flags. Also check the README from the repo they come from to get more details, if needed.

  1. Install labelImg. This is a Python package, which means you can install it via pip, but the one from GitHub is better. It saves annotations in the PASCAL VOC format.
  2. Annotate your dataset using labelImg.
  3. Use this script to convert the XML files generated by labelImg into a single CSV file.
  4. Use this script to separate the CSV file into two, one with training examples and one with evaluation examples. Let's call them train.csv and eval.csv. Images will be selected randomly and there are options to stratify examples by class, making sure that objects from all classes are present in both datasets. The usual proportions are 75% to 80% of the annotated objects used for training and the rest for the evaluation dataset.
  5. Create a "label map" for your classes. You can check some examples to understand what they look like. You can also generate one from your original CSV file with this script.
  6. Use this script to convert the two CSV files (train.csv and eval.csv) into two TFRecord files (eg. train.record and eval.record), a serialized data format that TensorFlow is most familiar with. You'll need the label map from the previous for this.

Choosing a neural network and preparing the training pipeline...

  1. Download one of the neural network models provided in this page. The ones trained in the COCO dataset are the best ones, since they were also trained on objects.
  2. Provide a training pipeline, which is a config file that usually comes in the tar.gz file downloaded in the previous step. If they don't come in the tar.gz, they can be found here. You can find a tutorial on how to create your own here.
    • The pipeline config file has some fields that must be adjusted before training is started. Its header describes which ones. Usually, they are the fields that point to the label map, the training and evaluation directories and the neural network checkpoint. In case you downloaded one of the models provided in this page, you should untar the tar.gz file and point the checkpoint path inside the pipeline config file to the "untarred" directory of the model (see this answer for help).
    • You should also check the number of classes. COCO has 90 classes, but your data set may have more or less.
    • There are additional parameters that may affect how much RAM is consumed by the training process, as well as the quality of the training. Things like the batch size or how many batches TensorFlow can prefetch and keep in memory may considerably increase the amount of RAM necessary, but I won't go over those here as there is too much trial and error in adjusting those.

Training the network

  1. Train the model. This is how you do it locally. Optional: in order to check training progress, TensorBoard can be started pointing its --logdir to the --model_dir of object_detection/model_main.py.

  2. Export the network, like this.

  3. Use the exported .pb in your object detector.

Tips

In the data augmentation section of the training pipeline, some options can be added or removed to try and make the training better. Some of the options are listed here.

@k1732036

This comment has been minimized.

Copy link

k1732036 commented Mar 28, 2018

It became easy for me to create my own object detector after reading your article.
Thank you!

@douglasrizzo

This comment has been minimized.

Copy link
Owner Author

douglasrizzo commented Jun 1, 2018

Hey, I'm glad you liked it. If you have any suggestions, leave a comment or maybe edit the article. Also, if anyone is interested, I just created a Python package that abstracts some of the more gruesome parts of using TensorFlow Object Detection API.

@freedylam

This comment has been minimized.

Copy link

freedylam commented Jun 7, 2018

Hi, douglasrizzo.

Do you mind list out how you put the folder of image, xml, csv, and TFRecord and config and how you issue the command?
I think I was missing in some part, when I doing training the train.py always get Killed.

already using CPU version of TF. and have 4GB memory of my ubuntu.

Thanks,
Freedy

@devajith

This comment has been minimized.

Copy link

devajith commented Sep 18, 2018

Can we do the splitting CSV by script ?

@douglasrizzo

This comment has been minimized.

Copy link
Owner Author

douglasrizzo commented Oct 10, 2018

@devajith you can, but it's a little tricky. You have to make sure you have examples of all classes in both the training and testing CSVs.

@Tsarpf

This comment has been minimized.

Copy link

Tsarpf commented Jan 8, 2019

The link to the package is dead. But thanks a lot for the cleaned up tutorial!

@douglasrizzo

This comment has been minimized.

Copy link
Owner Author

douglasrizzo commented Mar 18, 2019

Hey @Tsarpf, I just fixed the link to the package.

@douglasrizzo

This comment has been minimized.

Copy link
Owner Author

douglasrizzo commented Apr 23, 2019

Hi everyone, I've just updated the tutorial, mentioning a few scripts I created to help separate the main CSV into train and evaluation datasets, as well as generating the label map from the main CSV file. Since I've been the only one who tested these scripts, any issues or PRs in their repo are welcome.

@Rakshasv18

This comment has been minimized.

Copy link

Rakshasv18 commented May 2, 2019

@douglasrizzo hey ! getting error as No module named pandas! even when its installed

@douglasrizzo

This comment has been minimized.

Copy link
Owner Author

douglasrizzo commented May 2, 2019

@Rakshasv18 make sure you're using the Python version that has pandas installed. You may have Python 2 and 3 installed, or multiple versions of Python 3.

@Friso1987

This comment has been minimized.

Copy link

Friso1987 commented Jul 12, 2019

When I use the XML_to_CSV script my filenames are not included in the new CSV file. It fills them in as two dashes "--". How do I include my filename in convertion?

@douglasrizzo

This comment has been minimized.

Copy link
Owner Author

douglasrizzo commented Jul 12, 2019

@Friso1987 could you please check if your XML files contain the file names in them? Also, in which OS are you using the generate_csv.py script? I suspect that it won't work properly on Windows, since I've tested it only on Linux. You may submit an issue in the repo if you like, so I can take a look at it.

@Friso1987

This comment has been minimized.

Copy link

Friso1987 commented Jul 13, 2019

Found the mistake (stupid me). I got my XML files from an external source. When opening the XML files, I found out the filename was changed to "--". Probably for privacy reasons. I simply fixed it by saving the .xml files again in the LabelImg application, so it had a correct filename.

@MasudHaider

This comment has been minimized.

Copy link

MasudHaider commented Sep 24, 2019

Can this tutorial be run on google colaboratory?

@douglasrizzo

This comment has been minimized.

Copy link
Owner Author

douglasrizzo commented Oct 10, 2019

@MasudHaider I imagine so, as long as you are able to clone and install the TF Object Detection API into your Colab, the other steps are simple.

@Idakwo

This comment has been minimized.

Copy link

Idakwo commented Oct 28, 2019

@douglasrizzo Can you please share a file to recreate the environment you used? Like a requirement.txt or .yml. I am getting one error after the other after fixing a dependency error. There are several GitHub commits to TF's Object Detection repo correlating with different TF versions. I see changes in the repo to TF2.0. I'd be great to know which commit worked for you and the corresponding set up. Thanks

@douglasrizzo

This comment has been minimized.

Copy link
Owner Author

douglasrizzo commented Oct 29, 2019

Hi @ldakwo. I haven't updated this tutorial in a while. It works with tensorflow==1.13, but I can't pinpoint the version of the
tensorflow/models repo that is compatible with tf 1.13. I also know that object detection does not work with tf2.0 as of yet, so maybe that's one source for you errors.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.