Training custom YOLO v3 object detector

Python Lessons
Jan 23 · 3 min read

Keras implementation of YOLOv3 for custom detection:

Continuing from my previous tutorial, where I showed you how to prepare custom data for YOLO v3 object detection training, in this tutorial finally I will show you how to train that model.

First of all, I must mention that this code used in this tutorial originally is not my. I was inspired by qqwweee/keras-yolo3, he wrote this code to train custom YOLO v3 Keras model.

To make his code more efficient I made few modifications to it:
1. Originally used PIL library to load image for detection, I replaced it to OpenCV.
2. Originally used external font for drawing text on detected images — replaced with OpenCV.
3. Wrote my own detection scripts.

Quick Start:

1. Copy my GitHub repository. Must do this steps before.
2. Download YOLOv3 weights from YOLO website, or use wget command:

Copy downloaded weights file to model_data folder.
3. Convert the Darknet YOLO model to a Keras model:

4. Test YOLO v3 with image_detect.py or realtime_detect.py (modify used model and classes according to your needs)

Training guide:

There is 2 ways to train custom model:

  • train_bottleneck.py:
    Choose this method if you train on CPU, or want to train model faster (lower accuracy model). Required a lot of RAM and HDD space.
  • train.py:
    Choose this method if you train on GPU.

Modify before training:If using train.py:

Line 19. Change annotation_path to your file (learned to generate them in previous tutorial).
Line 20. Change log_dir, directory where to save trained model and checkpoints.
Line 21. Change classes_path to your classes file (learned to generate them in previous tutorial).
Line 22. anchors_path, don’t change this if you don’t know what you are doing.
Line 34. If training new model, leave it as it is “weights_path=’model_data/yolo_weights.h5'”, otherwise link your checkpoint.
Line 57. batch_size = 32, try to train with this, if you receive some kind of memory error, decrease this number.
Line 76. batch_size = 8, same as in Line 57. I tried to train with gtx1080ti, received memory error while using “batch_size = 32”, was able to train with 8.
Lines 63,64 82, 83. Increase epochs count for better model accuracy.

If using train_bottleneck.py:

Line 19. Change annotation_path to your file (learned to generate them in previous tutorial).
Line 20. Change log_dir, directory where to save trained model and checkpoints.
Line 21. Change classes_path to your classes file (learned to generate them in previous tutorial).
Line 22. anchors_path, don’t change this if you don’t know what you are doing.
Line 30. If training new model, leave it as it is “weights_path=’model_data/yolo_weights.h5'”, otherwise link your checkpoint.
Lines 72,73 86, 87, 105, 106. Increase epochs count for better model accuracy.

After finishing training in your logs file there should be created new “trained_weights_final.h5” model file. This will be used for custom detection.

Test trained model:

So I wrote two scripts to test YOLO v3 object detection and it’s accuracy:

  • Simpler one is test it with image_detect.py. Simply change “model_path” and “classes_path”, to your used files. Then at the end of file you can find “image = ‘hydrant.jpg’” line, here just change your test image name.
  • Second method I used in realtime_detect.py is that I try to grab my computer screen and detect objects on it. If you were folowing me for a while, you should remember that I found a way to grab screen faster using multiprocessing, so here I used same method to check YOLO v3 performance. If you were watching this tutorial you can see that with 1080x1080 screen on 1080ti gpu I am receiving around ~14 FPS. And with 800x600 I am receiving around ~20 FPS. This is much better that using object detection API !

Conclusion:

We successfully trained custom YOLO v3 object detection model on Keras. We are receiving quite nice performance results, and for you guys it should be much easier to train a new model than using object detection API, just download my code from my GitHub page. In next tutorial lets train this model to detect CS:GO enemies !

Related Tutorial List

Analytics Vidhya

Analytics Vidhya is a community of Analytics and Data Science professionals. We are building the next-gen data science ecosystem https://www.analyticsvidhya.com

Python Lessons

Written by

Analytics Vidhya

Analytics Vidhya is a community of Analytics and Data Science professionals. We are building the next-gen data science ecosystem https://www.analyticsvidhya.com

More From Medium

More from Analytics Vidhya

More from Analytics Vidhya

More from Analytics Vidhya

Get More Out of Google Colab

Welcome to a place where words matter. On Medium, smart voices and original ideas take center stage - with no ads in sight. Watch
Follow all the topics you care about, and we’ll deliver the best stories for you to your homepage and inbox. Explore
Get unlimited access to the best stories on Medium — and support writers while you’re at it. Just $5/month. Upgrade