Upload predictions

If for some reason you are not able to upload your trained YOLOv5 or YOLOv8 model to our platform, you can still benefit from model-assisted labeling by generating the predictions yourself and uploading them to our platform.

By uploading predictions to our platform, the predictions will become visible in our labeling interface. Similar to using our built-in label-assist tool, you will be able to verify and correct these predictions in the same way. The only difference is that you have to generate these predictions yourself.

Images to which you upload a prediction will still be marked as "unlabeled" as long as you have not manually inspected them and saved the annotations in our labeling interface.

To upload a prediction, it's as easy as setting the "prediction" attribute on a sample. We have provided conversion functions for both YOLOv5 and YOLOv8, to make it easy to directly upload their detections.

For completeness, we provide the following code snippets, which download the unlabeled samples, forward them through your model, and upload their detections.

Upload YOLOv5 predictions

To upload YOLOv5 predictions to the samples of your project, use the following (not-optimized) code snippet.

from trainyolo.client import Client, Project
from trainyolo.utils.yolov5 import format_boxes
from urllib.request import urlretrieve
import torch
import os

# initialize client
APIKEY = "YOUR_API_KEY"
client = Client(APIKEY)

# load unlabeled samples
PROJECT_NAME = "YOUR_PROJECT"
project = Project.get_by_name(client, PROJECT_NAME)
samples = project.get_samples(filter='UNLABELED')
  
# initialize YOLOv5 model
MODEL_PATH = "PATH_TO_YOLOV8_MODEL.pt"
SIZE, CONF, IOU = 640, 0.5, 0.45
model = torch.hub.load('ultralytics/yolov5', 'custom', path=model_path)
model.conf = CONF
model.iou = IOU

# create a temporary image location
image_loc = './images'
os.makedirs(image_loc, exist_ok=True)

for sample in samples:
  
  # download image
  image_path = os.path.join(image_loc, sample.name)
  if not os.path.exists(image_path):
      urlretrieve(sample.asset['url'], image_path)
  
  # forward image
  detections = model(image_path, size=SIZE)
  
  # upload detections
  boxes = detections.xyxy[0].cpu().numpy()
  sample.prediction = format_boxes(boxes)

Upload YOLOv8 predictions

from trainyolo.client import Client, Project
from trainyolo.utils.yolov8 import format_boxes, format_masks
from ultralytics import YOLO
from urllib.request import urlretrieve
import os

# initialize client
APIKEY = "YOUR_API_KEY"
client = Client(APIKEY)

# load unlabeled samples
PROJECT_NAME = "YOUR_PROJECT"
project = Project.get_by_name(client, PROJECT_NAME)
samples = project.get_samples(filter='UNLABELED')
  
# initialize YOLOv8 model
MODEL_PATH = "PATH_TO_YOLOV5_MODEL.pt"
SIZE, CONF, IOU = 640, 0.5, 0.45
model = YOLO(MODEL_PATH)

# create a temporary image location
image_loc = './images'
os.makedirs(image_loc, exist_ok=True)

for sample in samples:
  
  # download image
  image_path = os.path.join(image_loc, sample.name)
  if not os.path.exists(image_path):
      urlretrieve(sample.asset['url'], image_path)
  
  # forward image
  detections = model.predict(source=image_path, conf=CONF, iou=IOU, imgsz=SIZE)
  
  # OPTION 1: FOR DETECTION MODEL
  boxes, cls = detections[0].boxes.xyxy.cpu().numpy(), detections[0].boxes.cls.cpu().numpy()
  sample.prediction = format_boxes(boxes, cls)
  
  # OPTION 2: FOR SEGMENTATION MODEL
  masks, cls = detections[0].masks.cpu().numpy(), detections[0].boxes.cls.cpu().numpy()
  sample.prediction = format_masks(masks, cls)