Model Training

Overview

The first iteration of KB4YG uses public models namely COCO and YOLOv5. Theses models provide a good basis for the project but were trained using a broad dataset. Ultimately we will want to modify these models to be trained on our own dataset to improve recognition for specific objects.

Creating a dataset

We will need to gather our set of images and label our data. For this I recommend Roboflow with allows use to upload up 10k images for free and label them using their web app. After labeling our data we can export our dataset in several ML formats.

Training your Model

There are two main machine learning libraries TensorFlow and PyTorch. Since our models will be running on a Raspberry Pi we will be mostly using TensorFlow as it's the only one who natively supports the device. Although PyTorch can be used to train then export as TensorFlow lite model.

TensorFlow

In training\tensorflow here you'll a pair of Jupiter notebooks for training a TensorFlow Lite model. These can be used on your own computer by installing Jupiter Notebook or using Google Colab to train in the cloud. google-exp.ipynb is a training example provided by Google while roboflow-exp.ipynb is an example provide by roboflow for importing datasets from their website. Each guide should provide you with enough information to train your own model.

PyTorch

YOLOv5 provides a training guide for creating TFLite models base on the yolo dataset. You can find the guide in training\pytorch here.