YOLO builtin dataset

[nadinator]

🚀 The feature, motivation and pitch

A builtin torch.utils.data.DataSet for the YOLOv8 object detection model.
I've been working with this model for a while and I've had to use the ultralytics package for most of the work.
The problem is that this package is riddled with bugs, and trying to get help usually amounts to getting a GPT-generated reponse. So instead I've been going lower-level with the model and dataset that I have, which means that I've had to write my own custom dataset (extending the Dataset class). It would the life of any future programming a lot easier to not have to do this.

Alternatives

No response

Additional context

No response

cc @pmeier

[pmeier]

Hey @nadinator and thanks for the request. Could you detail what is special about YOLOv8 that it needs a custom dataset? We have builtin datasets that support object detection, e.g. torchvision.datasets.CocoDetection or torchvision.datasets.VOCDetection.

[nadinator]

Hi @pmeier and thanks for the quick reply! YOLO has a different format for the image/annotation pairs than those other datasets. You can read more about it here.

[pmeier]

If all the information is there and you just need to change some formatting to work with an arbitrary model, your best bet is to write a custom transform that handles this for you.

Using our newly introduced v2 transforms, this is fairly straightfoward. For example, the conversion to normalized XYWH bounding boxes looks like

import torch
from torchvision import datapoints, datasets
from torchvision.transforms import v2
from torchvision.transforms.v2 import functional as F


class CocoDetection2Yolov2(v2.Transform):
    def forward(self, image, target):
        height, width = F.get_size(image)

        xyxy_boxes = target["boxes"]
        xywh_boxes = F.convert_bounding_box_format(xyxy_boxes, new_format=datapoints.BoundingBoxFormat.XYWH)
        normalized_xywh_boxes = xywh_boxes / torch.tensor([width, height, width, height])

        ...


dataset = datasets.CocoDetection(
    ...,
    transforms=CocoDetection2Yolov2(),
)
dataset = datasets.wrap_dataset_for_transforms_v2(dataset, target_keys="all")

[nadinator]

I actually need to work with the Yolov8n model. Would this approach still work @pmeier ?

[pmeier]

@nadinator We can't provide support for third-party models. According to the link in #7861 (comment) the COCO dataset is supported. Thus, you'll have to implement a transform or the like to make our dataset compatible with your model.

[oke-aditya]

I think the format is (cx, cy, w, h) Center of box is considered instead of corner for YOLO Type of models.
I guess we do have format conversion for the same. You can then normalize the box

import torch
from torchvision import datapoints, datasets
from torchvision.transforms import v2
from torchvision.transforms.v2 import functional as F


class CocoDetection2Yolov2(v2.Transform):
    def forward(self, image, target):
        height, width = F.get_size(image)

        xyxy_boxes = target["boxes"]
        xywh_boxes = F.convert_bounding_box_format(xyxy_boxes, new_format=datapoints.BoundingBoxFormat.CXCYWH)
        normalized_xywh_boxes = xywh_boxes / torch.tensor([width, height, width, height])

        ...


dataset = datasets.CocoDetection(
    ...,
    transforms=CocoDetection2Yolov2(),
)
dataset = datasets.wrap_dataset_for_transforms_v2(dataset, target_keys="all")