How to Train Yolov8 Object Detection on Custom Dataset

A symbolic link is needed as Google Colab is not good at handling space like in “My Drive”. You can use the symbolic link “/content/yoloV8_fruit_veggies_detection” to access files in the directory

Note: You need to create the directory named yoloV8_fruit_veggies_detection in your Google Drive

!ln -s "/content/gdrive/My Drive/yoloV8_fruit_veggies_detection" /content/yoloV8_fruit_veggies_detection


Just in case u want to remove a symbolic link

!unlink /content/yoloV8_fruit_veggies_detection



!nvidia-smi


Get the Current working dir & set the HOME variable. The HOME variable will be used to access files stored locally on Colab temporary allocated storage space

import os
HOME = os.getcwd()
print(HOME)


Install ultralytics which is the package containing YOLOv8 & Check if YOLOv8 is working as expected

!pip install ultralytics==8.0.20

from IPython import display
display.clear_output()

import ultralytics
ultralytics.checks()


Configuration file is where you define all parameters which model is going use for training. Examples of these are:

No. of epochs

batch size

optimizer …


!cat /usr/local/lib/python3.9/dist-packages/ultralytics/yolo/cfg/default.yaml


Import Yolo from the ultralytics package.


from ultralytics import YOLO

from IPython.display import display, Image


Yolo CLI interface is the best way to get started with training, validating or inferencing i.e., if you dont have to make any modifications to the code. If you want more information you can goto Ultralytics YOLO Docs


yolo task=detect mode=train model=yolov8n.yaml args...
classify predict yolov8n-cls.yaml args...
segment val yolov8n-seg.yaml args...
export yolov8n.pt format=onnx args...


Yolo mode=predict automatically downloads models from latest YOLOv8 release & saved the result of prediction to runs/predict.

%cd {HOME}
!yolo task=detect mode=predict model=yolov8n.pt conf=0.25 source='https://media.roboflow.com/notebooks/examples/dog.jpeg' save=True


Displays the predicted image in the notebook along with BBoxes

%cd {HOME}
Image(filename='runs/detect/predict/dog.jpeg', height=600)


We’ll be downloading a open source dataset from Roboflow to train the model. if you want to inspect the dataset you can download it from this link: https://universe.roboflow.com/bohni-tech/fruits-and-vegi

!mkdir {HOME}/datasets
%cd {HOME}/datasets

!pip install roboflow --quiet

from roboflow import Roboflow
rf = Roboflow(api_key="REPLACE_WITH_UR_API_FROM_ROBOFLOW")
project = rf.workspace("bohni-tech").project("fruits-and-vegi")
dataset = project.version(13).download("yolov8")


This is to start executing from Gdrive dir, since model is saved in dir where execution starts. There is presently no way to specify a custom path to a directory to save the trained YoloV8 model. So, best method is to start model execution from the GDrive in which you want the model to be saved

I have created a subdirectory named train_march_23_2023 in my Google Drive in which i intend to save the model. You change the name & create your own sub directory

%cd /content/yoloV8_fruit_veggies_detection/train_march_23_2023
!pwd


Starts training the model, the details of parameters are as below:

task=detect – Trains the model for object detection

model=yolov8s.pt – Model to train is yolov8s which is a small sized model, this can be easily depolyed on edge device<

data={dataset.location}/data.yaml – Location of the data.yaml file which in turn has the dataset location

epochs=3 – Train for 3 epochs

imgsz=1280 – Size of the training images, in this case its 1280×1280

plots=True – Plots of loss, accuracy, Precision, Recall etc

batch=8 – No of images per batch


!yolo task=detect mode=train model=yolov8s.pt data={dataset.location}/data.yaml epochs=3 imgsz=1280 plots=True batch=8


/runs/detect/train7/weights/best.pt – You will have to change this path to point to the latest directory, each training run will generate a seperate directory i.e., train7 – this will be train8 on next run


resume_location='/content/yoloV8_fruit_veggies_detection/train_march_23_2023/runs/detect/train7/weights/best.pt'


resume=True – Only works if training has been stopped midway for example: If you had started training with epoch=30 & training stopped at epoch=20 , then training will resume from epoch=20

model={resume_location} – Location of the checkpoint


!yolo task=detect mode=train model=yolov8s.pt data={dataset.location}/data.yaml epochs=30 imgsz=1280 plots=True batch=8 resume=True model={resume_location}


Check how accurately the model is predicting with Confusion Matrix

Path – Is defined below to Google Drive where the Confusion Matrix is saved


gdrive_path = '/content/yoloV8_fruit_veggies_detection/train_march_23_2023/runs/detect'

Image(filename=f'{gdrive_path}/train5/confusion_matrix.png', width=600)


The image gives the Accuracy measures as listed below, these are used to judge how well the model has been trained

Training Loss

Validation Loss

mAP50-95

Precision Curve

Recall Curve

Image(filename=f'{gdrive_path}/train5/results.png', width=600)

Check out how well the model is trained by predicting on a test image

Image(filename=f'{gdrive_path}/train5/val_batch2_pred.jpg', width=600)


Check how model performs on validation data for different classes for following parameters:

mAP50

mAP50-95

P

R


!yolo task=detect mode=val model={gdrive_path}/train5/weights/best.pt data={dataset.location}/data.yaml


Use the model to inference on the test images & have a look at the inferenced images

!yolo task=detect mode=predict model={gdrive_path}/train5/weights/best.pt conf=0.25 source={dataset.location}/test/images save=True

import glob
from IPython.display import Image, display

for image_path in glob.glob(f'{gdrive_path}/predict/*.jpg’)[3:8]:
display(Image(filename=image_path, width=600))
print(“\n”)