supervision icon indicating copy to clipboard operation
supervision copied to clipboard

Create new cookbook for utilizing supervision methods to easily create YOLO datasets for training

Open xaristeidou opened this issue 7 months ago • 7 comments

Search before asking

  • [X] I have searched the Supervision issues and found no similar feature requests.

Description

I find myself creating dataset structures and split in train, valid, test and images, labels folders multiple times. The whole process could easily be automated.

Use case

There are methods that currently exist to load a dataset sv.DetectionDataset.from_yolo(), split in selected ratio sv.DetectionDataset.split() and export to YOLO format sv.DetectionDataset.to_yolo().

Nevertheless, in creation of YOLO model training dataset structure, someone must write a custom split in train/valid/test (because split() is able to split only in two parts), and also create manually the train/valid/test folders needed for the preparation of the dataset. (As of my knowledge ultralytics YOLO models require by default to have train/valid folders that contain valid and not empty annotations, test folder can be empty).

For that reason I propose a new method to be added in sv.DetectionDataset which will combine the arguments of from_yolo(), split(), to_yolo() and will run the whole backend for creating train/valid/test folder and images/labels subfolders along with data.yaml file.

At this current point I have developed an implementation of such method which provides the ability to the user to create a YOLO dataset structure with a single line of code. An example of executing such a process can be seen in the following example:

import supervision as sv

dataset_directory = "/path/to/directory"

sv.DetectionDataset.create_yolo_dataset(
    images_directory_path=f"{dataset_directory}/images",
    annotations_directory_path=f"{dataset_directory}/labels",
    data_yaml_path=f"{dataset_directory}/data.yaml",
    train_ratio=0.7,
    valid_ratio=0.15,
    folders_export_path=f"{dataset_directory}",
    data_yaml_export_path = f"{dataset_directory}/data.yaml",
)

Additional

  • The test_ratio is automatically calculated based on the train_ratio, valid_ratio.
  • The arguments provided in the example are all mandatory arguments to be passed.
  • The user can provide any additional optional argument incorporated in from_yolo(), split(), to_yolo().

Let me know if you like this idea, and if you want to submit a PR with the initial implementation.

Are you willing to submit a PR?

  • [X] Yes I'd like to help by submitting a PR!

xaristeidou avatar Jul 21 '24 16:07 xaristeidou