doutdex / yolov2-tensorflow-2.0

Just another YOLO V2 implementation. Train your own dataset in a jupyter notebook!

YOLO V2 with TensorFlow 2.0

Here is a jupyter notebook featuring a complete implementation from scratch of YOLOV2 with TensorFlow 2.0 :

• Dataset pipeline with data augmentation
• Training from YOLO pretrained weights
• Visualization of object detection

I use this notebook to train a model to detect crop and weeds in a field. The goal is to detect crop in real time for tractor guidance and to detect weeds to remove them.

Original paper : YOLO9000: Better, Faster, Stronger by Joseph Redmond and Ali Farhadi.

Files

• Yolo_V2_tf_2.ipynb : Yolo V2 implementation with Tensorflow 2.0
• Yolo_V2_tf_eager.ipynb : old notebook, Yolo V2 implementation with Tensorflow 1.x with eager execution

Requirements

• tensorflow 2.0
• imgaug
• cv2

Before using the notebook

• Download pretrained weights here. Place this weights file in notebook directory and name it yolo.weights
• The training requires four directories containing images and annotations :

train_image_folder/ : contains images files used during training (png format)

train_annot_folder/ : contains annotations in PASCAL VOC format (one xml file for each image)

val_image_folder/ : contains images files used for validation

val_annot_folder/ : contains annotations in PASCAL VOC format

Using the notebook

• Define object's labels to detect (same labels as defined in PASCAL VOC xml file). Example :

	LABELS           = ('sugarbeet', 'weed')

• Define image size in dataset and YOLO grid size. Image size must be YOLO grid size * 32.

	IMAGE_H, IMAGE_W = 512, 512
	GRID_H,  GRID_W  = 16, 16 # GRID size = IMAGE size / 32

• Define train batch size and validation batch size : depends on image size and video card RAM.

	TRAIN_BATCH_SIZE = 10
	VAL_BATCH_SIZE   = 10

• Define path to dataset directories.

	# Train and validation directories

	train_image_folder = 'data/train/image/'
	train_annot_folder = 'data/train/annotation/'
	val_image_folder = 'data/val/image/'
	val_annot_folder = 'data/val/annotation/'

That's it, just run notebook cells to train YOLO on your own data!

Example of use

YOLO model trained on sugarbeet and weed dataset (two labels) :

Credits

Many thanks to these great repositories:

https://github.com/experiencor/keras-yolo2

https://github.com/allanzelener/YAD2K

and to this very good explanation of the YOLO V2 loss function:

https://fairyonice.github.io/Part_4_Object_Detection_with_Yolo_using_VOC_2012_data_loss.html