Skip to content

Ikomia-hub/train_yolo_v5

Repository files navigation

Algorithm icon

train_yolo_v5


Stars Website GitHub
Discord community

Train Ultralytics YoloV5 object detection models.

Desk object detection

🚀 Use with Ikomia API

1. Install Ikomia API

We strongly recommend using a virtual environment. If you're not sure where to start, we offer a tutorial here.

pip install ikomia

2. Create your workflow

from ikomia.dataprocess.workflow import Workflow

# Init your workflow
wf = Workflow()    

# Add dataset loader
coco = wf.add_task(name="dataset_coco")

coco.set_parameters({
    "json_file": "path/to/json/annotation/file",
    "image_folder": "path/to/image/folder",
    "task": "detection",
}) 

# Add training algorithm
train = wf.add_task(name="train_yolo_v5", auto_connect=True)
train.set_parameters({"dataset_folder": "path/to/where/will/be/saved/yolov5/format/dataset"}) 

# Launch your training on your data
wf.run()

☀️ Use with Ikomia Studio

Ikomia Studio offers a friendly UI with the same features as the API.

  • If you haven't started using Ikomia Studio yet, download and install it from this page.

  • For additional guidance on getting started with Ikomia Studio, check out this blog post.

📝 Set algorithm parameters

  • model_name (str) - default 'yolov5s': Name of the pre-trained model. Additional models available:

    • yolov5n
    • yolov5m
    • yolov5l
    • yolov5x
  • dataset_folder (str): Path to the re-structured dataset to YOLOv5 format will be saved.

  • input_width (int) - default '512': Width of the input image.

  • input_height (int) - default '512': Height of the input image.

  • epochs (int) - default '10': Number of complete passes through the training dataset.

  • batch_size (int) - default '16': Number of samples processed before the model is updated.

  • dataset_split_ratio (float) – default '0.9': Divide the dataset into train and evaluation sets ]0, 1[.

  • output_folder (str, optional): path to where the model will be saved.

Parameters should be in strings format when added to the dictionary.

from ikomia.dataprocess.workflow import Workflow

# Init your workflow
wf = Workflow()    

# Add dataset loader
coco = wf.add_task(name="dataset_coco")

coco.set_parameters({
    "json_file": "path/to/json/annotation/file",
    "image_folder": "path/to/image/folder",
    "task": "detection",
}) 

# Add training algorithm
train = wf.add_task(name="train_yolo_v5", auto_connect=True)
train.set_parameters({
    "model_name": "yolov5s",
    "dataset_folder": "path/to/where/will/be/saved/yolov5/format/dataset",
    "epochs": "5",
    "batch_size": "2",
    "input_width": "512",
    "input_height": "512",
    "dataset_split_ratio": "0.9"
}) 

# Launch your training on your data
wf.run()