Datasets

How to add a new dataset

This guide explains how to add a new dataset to WIBE framework. For more examples, refer to the wibench.datasets module.

Create your_dataset.py file in user_plugins directory.

Here we have an example for image based dataset.

from wibench.datasets import BaseDataset
from wibench.typing import ImageObject

class MyDataset(BaseDataset):

    def __init__(self, parametrs_of_dataset):
        ...
        # Any initialization dataset may need

    def __len__(self) -> int:
        # Length of dataset if available for progress bar.

    def generator(self) -> Generator[ImageObject, None, None]:
            # Yields images from directory.
            ...
            yield ImageObject(image_id, torch_image)

If it is possible to get number of samples in dataset, you may inherit from RangeBaseDataset.

from wibench.datasets import RangeBaseDataset
from wibench.typing import ImageObject

class MyDataset(RangeBaseDataset):

    # Pipeline type depends on returned data type
    pipeline_type = PipelineType.IMAGE

    def __init__(self, parametrs_of_dataset, sample_range: Optional[Tuple[int, int]] = None):
        ...
        super().__init__(sample_range, self.__len__())
        ...

    def __len__(self) -> int:
        ...

    def generator(self) -> Generator[ImageObject, None, None]:
            # Yields images from directory. Alternatively, PromptObject
            # may be returned (pipeline_type should be changed accordingly)
            ...
            yield ImageObject(image_id, torch_image)

Implemented datasets

class wibench.datasets.base.ImageFolderDataset(path: Union[Path, str], preload: bool = False, img_ext: List[str] = ['png', 'jpg'], sample_range: Optional[Tuple[int, int]] = None)

Concrete dataset implementation loading images from a directory.

Supports common image formats with optional preloading.

Parameters

pathUnion[Path, str]

Directory path containing images

preloadbool

Whether to load all images into memory upfront

img_extList[str]

Image file extensions to include (default: [‘png’, ‘jpg’])

sample_rangeOptional[Tuple[int, int]]

Optional (start, end) index range to subset the dataset (including both borders)

class wibench.datasets.base.PromptFolderDataset(path: Union[Path, str], prompt_ext: List[str] = ['txt', 'csv'], sample_range: Optional[Tuple[int, int]] = None, separator: str = '\n')

Concrete dataset implementation loading prompts from a directory. Directory should contain a number of “.txt” or “.csv” files with prompts, in one file prompts are separated by separator. All prompts are preloaded.

Parameters

pathUnion[Path, str]

Directory path containing images

prompt_extList[str]

File extensions to include (default: [‘txt’, ‘csv’])

sample_rangeOptional[Tuple[int, int]]

Optional (start, end) index range to subset the dataset (including both borders). Default: None (full dataset)

separatorstr

Separator for prompts in one file, default is line break

class wibench.datasets.mscoco.mscoco.MSCOCO(split: str = 'val', sample_range: Optional[Tuple[int, int]] = None, cache_dir: Optional[str] = None, return_prompt: bool = False, timeout: int = 3600)

Dataset loader for MS-COCO (Common Objects in Context) images.

Provides access to the COCO 2017 dataset images through HuggingFace Datasets, supporting both validation and training splits with optional caching.

Parameters

splitstr

Dataset split to load (‘train’ or ‘val’)

sample_rangeOptional[Tuple[int, int]]

Optional (start, end) index range to subset the dataset

cache_dirOptional[str]

Directory to cache downloaded dataset files

return_promptbool

If enabled, returns image captions instead of images (default False)

timeoutint

Timeout for dataset download