Training custom YOLO v3 object detector

Posted August 26, 2019 by Rokas Balsys



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:

wget https://pjreddie.com/media/files/yolov3.weights
Copy downloaded weights file to model_data folder.
3. Convert the Darknet YOLO model to a Keras model:
python convert.py model_data/yolov3.cfg model_data/yolov3.weights model_data/yolo_weights.h5
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 !
Conslusion:

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 easer 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 !