How to assign category names to a dataset created with ImageFolder?

I created an image dataset (for object identification) following the ImageFolder approach. The created dataset only has the class ids, not their names anywhere.

The documentation on load_dataset doesn’t explain well how I would add this (or any other extra metadata).

Hi! You can define labels in the load_dataset call as follows:

import datasets

class_names = [...]
features = datasets.Features(
    {
        "image": datasets.Image(),
        "objects": {
            "bbox": datasets.Sequence(datasets.Value("float32"), length=4), 
            "categories": datasets.ClassLabel(names=class_names)
        }
    }
)
ds = datasets.load_dataset(..., features=features)
1 Like