Training Your Own AI Image Recognition Model

Until now you have used pre-trained models that recognize common objects. Today you train your own, so the Pi can recognize something specific: your logo, a water bottle, a banana, even your face. You will learn how models actually learn from labeled images, how to collect and label a dataset, how to train with YOLO, and how to run your custom model on the Pi. Understanding this is what turns AI from a black box into a tool you control.

How recognition is learned

An AI model does not know what a banana is; it learns from examples. You give it many images, each labeled with what it contains and where, and during training the model gradually learns the visual patterns and features that distinguish that object. Afterward it can recognize the object in new photos it has never seen. The single most important factor is the dataset: a large, varied, well-labeled set produces a strong model, and a small or sloppy one produces a weak model. Good data matters more than anything else.

Collecting a good dataset

Gather images of your object from the Pi camera, a phone, or a licensed online source. Aim for many, at least ten per category and ideally far more, captured from different angles, in different lighting, and against different backgrounds, sometimes with other objects nearby. Variety is what teaches the model to recognize the object in the real world rather than memorizing one photo. Data augmentation, automatically rotating or adjusting images, helps stretch a small set. For common objects, check for a ready-made open dataset first before shooting your own.

Labeling and the dataset structure

Once collected, you label each image, drawing a box around the object and naming it, using a tool like Roboflow or LabelImg. For YOLO the exported dataset has a specific structure: a data file listing the label names and folder paths, and train and validation folders, each holding the images and matching label files that record the box coordinates. The labels are how the model knows where the object is, since it has no prior knowledge. Getting the folder paths right in the data file is a common source of errors.

Why split train and validation

You never use all your data for training. You hold back a portion as a validation set, images the model does not train on, so you can honestly measure how well it does on things it has not seen. Testing a model only on data it trained on tells you nothing, because it could simply have memorized those. The validation set is the unseen exam that reveals whether the model actually learned to generalize.

Training, epochs, and choosing the weights

Training runs the images through the model repeatedly. Two settings matter most: the image size, where bigger means higher quality but more computation, and the number of epochs, how many times it iterates over the data, often 10 to 30 depending on the dataset. Training is heavy, so it is fine, and often faster, to train on your PC or the cloud and copy the finished model to the Pi. Training produces two weight files: the best, from the epoch that performed best, and the last, from the final epoch. Prefer the best, because performance can actually dip in later epochs.

Working through it

Collect images. Shoot at least 20-plus images of your chosen object from varied angles, lighting, and backgrounds, or download a licensed dataset for a common object.

Label the data. Use Roboflow or LabelImg to draw a box around the object in each image and name it, then export in YOLO format with the data file and train/validation folders.

Point the data file at your folders. Confirm the paths in the data file match where your images and labels actually live; a wrong path is a frequent training error.

Train the model. Install the YOLO dependencies, then run training with a base model, your data file, an image size, and a number of epochs. Train on a PC or cloud if the Pi is slow, and wait for it to finish.

Run your custom model. Take the best weights file, copy it to the Pi, and run detection just like lesson 2 but pointing at your model. Test it on an image or the live camera and watch it box your object.

The YOLO dataset data file

# data.yaml describes your dataset
path: /home/pi/documents/academy/month3/lesson7   # dataset root
train: train/images        # training images
val: valid/images          # validation (unseen) images

names:
  0: banana                # your label(s); rename freely

# folder layout expected by YOLO:
#   train/images/*.jpg  train/labels/*.txt
#   valid/images/*.jpg  valid/labels/*.txt
# each label .txt lists the object's class and box coordinates