Train YOLO with DarkFlow

Retrain YOLO using DarkFlow

DarkFlow is the TensorFlow specific implementation of the DarkNet. In this tutorial we will use this framework to retrain a tiny-yolo model for two classes. More information about the DarkFlow can be found on the official site here. As usual, we prepared a demo for training YOLO with DarkFlow in Google Colaboratory Notebook: https://drive.google.com/file/d/1CYoSezen36kP9NLNJc42ycNztIgxS25t/view?usp=sharing.

Let's begin with installation of additional softwares like: CUDA, cuDNN, python-opencv etc. If you have these dependencies already installed, then you can jump to the Install DarkFlow part; It is optional, but recommended to install nVidia CUDA and cuDNN libraries before you start train your model. A short description about the installation can be found here: Install nVidia CUDA and cuDNN.

Install TensorFlow

There are a couple way to install the TensorFlow, you can find them here. Right now we will use the pip python package manager to install it.

Install TensorFlow for Python 2.x

• $ pip install tensorflow # Python 2.7; CPU support (no GPU support)
• $ pip install tensorflow-gpu # Python 2.7; GPU support

Install TensorFlow for Python 3.x

• $ pip3 install tensorflow # Python 3.n; CPU support (no GPU support)
• $ pip3 install tensorflow-gpu # Python 3.n; GPU support

Install Numpy, OpenCV, Cython

• $ apt-get update
• $ pip3 install numpy
• $ apt-get install python-opencv -y
• $ pip install cython

Install DarkFlow

Finally, we have all the dependencies installed so, we can install the DarkFlow right now.

• Clone the repository with: https://github.com/thtrieu/darkflow.git;
• Run the python3 setup.py build_ext --inplace command from the target folder;
• Run pip install .

Train with DarkFlow

At this point we are ready to retrain our tiny-yolo model for two classes (Winnie-the-pooh and Tiger) . For this reason we will use the dataset generater from the previous article.

Clone the dataset generator from the repo, install it and generate some learning data.

• $ git clone https://github.com/szaza/dataset-generator.git
• $ cd dataset-generator
• $ git checkout feature/colaboratory-tutorial // this branch contains the demo data for this example
• $ ./gradlew clean build -xtest
• $ ./gradlew run

Configure DarkFlow by modifying the configuration file and labels.txt file. Make a copy from cfg/tiny-yolo-voc.cfg and create a cfg/tiny-yolo-voc-2c.cfg file with the same content. Change the line 114 to filters=35 and the line 120 and set classes=2.

Copy the labels.txt file from the dataset-generator/dataset/output/VOCLabels and replace the original one from the DarkFlows root directory.

It can take days to train YOLO from scratch. A faster option is to train it by using some pretrained weights. In this tutorial we are using the tiny-yolo-voc.weights file, what can be downloaded from here: https://oc.codespring.ro/s/Jgyo6N4Jen3ma2P/download; Download it and place it to the darkflow/weights directory and rename it to tiny-yolo-voc.cfg.

Start training with the following command:

$ python3 flow --model ./cfg/tiny-yolo-voc-2c.cfg --train --dataset "<path to the generator>dataset-generator/dataset/output/VOCDataSet/JPEGImages" --annotation "<path to the generator>/dataset-generator/dataset/output/VOCDataSet/Annotations" --gpu 0.8 --epoch 10000 --load ./weights/tiny-yolo-voc.weights

The outputs in the console should look like on the picture below:

The training process can be stopped at any moment by pressing CTR + C command in the terminal window. The recommended number of steps for this demo dataset is around 10000-15000.

To save into protobuff file use the ./flow --model cfg/tiny-yolo-voc-2c.cfg --load -1 --savepb command.

The newly trained weights can be saved into a protobuff file with the following command:

• $ ./flow --model cfg/tiny-yolo-voc-2c.cfg --load -1 --savepb

The tensorflow java example project can be used to test the protobuff file.