mmlearn.datasets.imagenet.ImageNet

class ImageNet(root_dir, split='train', transform=None, target_transform=None, mask_generator=None)[source]

Bases: ImageFolder

ImageNet dataset.

This is a wrapper around the ImageFolder class that returns an Example object.

Parameters:

  • root_dir (str) – Path to the root directory of the dataset.

  • split ({"train", "val"}, default="train") – The split of the dataset to use.

  • transform (Optional[Callable], optional, default=None) – A callable that takes in a PIL image and returns a transformed version of the image as a PyTorch tensor.

  • target_transform (Optional[Callable], optional, default=None) – A callable that takes in the target and transforms it.

  • mask_generator (Optional[Callable], optional, default=None) – A callable that generates a mask for the image.

Methods

Attributes

__getitem__(index)[source]

Get an example at the given index.

Return type:

Example

property id2label: dict[int, str]

Return the label mapping.

property zero_shot_prompt_templates: list[str]

Return the zero-shot prompt templates.