Skip to content

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

License

Notifications You must be signed in to change notification settings

jmpap/YOLOV2-Tensorflow-2.0

Repository files navigation

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

About

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

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published