7. Training an Object Detection Model with YOLOv8

In this tutorial we will download a public model from the AI library to inference and train on custom data locally.Here we will use a YOLOv8 model.

Create a venv and install the requirements for the yolov8 package hereThen, import the modules required for the scripts in this tutorial.

# !pip install torch torchvision imgaug "scikit-image<0.18"
import matplotlib.pyplot as plt
from PIL import Image
import numpy as np
import dtlpy as dl
import json
from dtlpy.ml import train_utils

7.1. Create a Project and a Dataset

We will use a public fruits dataset. We create a project and a dataset and upload the data with 3 labels of fruit.NOTE: You might need to change the location of the items, which currently points to the root of the documentation repository. If you downloaded the dtlpy documentation repo locally, this should work as-is.

project = dl.projects.create('Fruit - Model Mgmt')
dataset = project.datasets.create('Fruit')
dataset.to_df()
_ = dataset.items.upload(local_path='../../../../assets/sample_datasets/FruitImage/items/*',
                         local_annotations_path='../../../../assets/sample_datasets/FruitImage/json')
dataset.add_labels(label_list=['orange', 'banana', 'apple'])
dataset.open_in_web()

Now we’ll add the train and validation sets to the dataset metadata:

subsets = {'train': json.dumps(dl.Filters(field='dir', values='/train').prepare()),
           'validation': json.dumps(dl.Filters(field='dir', values='/validation').prepare())}
dataset.metadata['system']['subsets'] = subsets
dataset.update(True)

7.2. Clone the Public Model Into Your Project

We’ll get and clone the public yolo pretrained model (you can view the public models in the public Dataloop Github).You can view all publicly available models by using a Filter. Here we will use a YOLOv8 model pretrained on the COCO dataset.

import dtlpy as dl
filters = dl.Filters(resource=dl.FiltersResource.MODEL, use_defaults=False)
filters.add(field='scope', values='public')
dl.models.list(filters=filters).to_df()
# get the public model
pretrained_model = dl.models.get(model_name='pretrained-yolo-v8')
model = pretrained_model.clone(model_name='fruits-model',
                               dataset=dataset,
                               project_id=project.id,
                               configuration={'epochs': 10,
                                              'batch_size': 4,
                                              'imgz': 640,
                                              'id_to_label_map': {(v - 1): k for k, v in
                                                                  dataset.instance_map.items()}
                                              })

7.3. Train on Your Dataset

We’ll load the new, untrained model into the adapter and prepare the local dataset to be used for training.

adapter = dl.packages.build(package=model.package,
                            init_inputs={'model_entity': model},
                            module_name='model-adapter')

7.4. Start the Training

The package, model, and data are now prepared. We are ready to train!

print("Training {!r} on dataset {!r}".format(model.name, dataset.name))
adapter.train_model(model=model)

We can list all Artifacts associated with this Package, and add more files that are needed to load or run the model.

adapter.model.artifacts.list_content()

7.5. Predict Your Newly Trained Model

With everything in place, we will load our model and view an item’s prediction.

item = dl.items.get(item_id='6110d4a41467ded7a8c2a23d')
annotations = adapter.predict_items([item], with_upload=True)
image = Image.open(item.download())
plt.imshow(item.annotations.show(image,
                                 thickness=5))
print('Classes found: {}'.format([ann.label for ann in annotations[0]]))