lightly.data

The lightly.data module provides a dataset wrapper and collate functions.

.dataset

class lightly.data.dataset.LightlyDataset(input_dir: Optional[str], transform: Compose = None, index_to_filename: Callable[[VisionDataset, int], str] = None, filenames: List[str] = None, tqdm_args: Dict[str, Any] = None, num_workers_video_frame_counting: int = 0)

Provides a uniform data interface for the embedding models.

Should be used for all models and functions in the lightly package. Returns a tuple (sample, target, fname) when accessed using __getitem__.

The LightlyDataset supports different input sources. You can use it on a folder of images. You can also use it on a folder with subfolders with images (ImageNet style). If the input_dir has subfolders, each subfolder gets its own target label. You can also work with videos (requires pyav). If there are multiple videos in the input_dir each video gets a different target label assigned. If input_dir contains images and videos only the videos are used.

Can also be used in combination with the from_torch_dataset method to load a dataset offered by torchvision (e.g. cifar10).

Parameters

input_dir – Path to directory holding the images or videos to load.
transform – Image transforms (as in torchvision).
index_to_filename – Function which takes the dataset and index as input and returns the filename of the file at the index. If None, uses default.
filenames – If not None, it filters the dataset in the input directory by the given filenames.

Examples

>>> # load a dataset consisting of images from a local folder
>>> # mydata/
>>> # `- img1.png
>>> # `- img2.png
>>> # `- ...
>>> import lightly.data as data
>>> dataset = data.LightlyDataset(input_dir="path/to/mydata/")
>>> sample, target, fname = dataset[0]
>>>
>>> # also works with subfolders
>>> # mydata/
>>> # `- subfolder1
>>> #     `- img1.png
>>> # `- subfolder2
>>> # ...
>>>
>>> # also works with videos
>>> # mydata/
>>> # `- video1.mp4
>>> # `- video2.mp4
>>> # `- ...

dump(output_dir: str, filenames: Optional[List[str]] = None, format: Optional[str] = None)

Saves images in the dataset to the output directory.

Will copy the images from the input directory to the output directory if possible. If not (e.g. for VideoDatasets), will load the images and then save them to the output directory with the specified format.

Parameters

output_dir – Output directory where the image is stored.
filenames – Filenames of the images to store. If None, stores all images.
format – Image format. Can be any pillow image format (png, jpg, …). By default we try to use the same format as the input data. If not possible (e.g. for videos) we dump the image as a png image to prevent compression artifacts.

classmethod from_torch_dataset(dataset, transform=None, index_to_filename=None)

Builds a LightlyDataset from a PyTorch (or torchvision) dataset.

Parameters

dataset – PyTorch/torchvision dataset.
transform – Image transforms (as in torchvision).
index_to_filename – Function which takes the dataset and index as input and returns the filename of the file at the index. If None, uses default.

Returns

A LightlyDataset object.

Examples

>>> # load cifar10 from torchvision
>>> import torchvision
>>> import lightly.data as data
>>> base = torchvision.datasets.CIFAR10(root="./")
>>> dataset = data.LightlyDataset.from_torch_dataset(base)

get_filenames() → List[str]: Returns all filenames in the dataset.

get_filepath_from_filename(filename: str, image: <module 'PIL.Image' from '/datasets/actions-runner/core_gpu_runner_01/_work/lightly-core/lightly-core/.venv/lib/python3.12/site-packages/PIL/Image.py'> = None)

Returns the filepath given the filename of the image

There are three cases:

The dataset is a regular dataset with the images in the input dir.
The dataset is a video dataset, thus the images have to be saved in a temporary folder.
The dataset is a torch dataset, thus the images have to be saved in a temporary folder.

Parameters

filename – The filename of the image
image – The image corresponding to the filename

Returns

The filename to the image, either the existing one (case 1) or a newly created jpg (case 2, 3)

property transform: Getter for the transform of the dataset.

.multi_view_collate

class lightly.data.multi_view_collate.MultiViewCollate

Collate function that combines views from multiple images into a batch.

This collate function processes a batch of tuples, where each tuple contains multiple views of an image, a label, and a filename. It outputs these as separate grouped tensors for easy batch processing.

Example

>>> transform = SimCLRTransform()
>>> dataset = LightlyDataset(input_dir, transform=transform)
>>> dataloader = DataLoader(
...     dataset, batch_size=4, collate_fn=MultiViewCollate()
... )
>>> for views, targets, filenames in dataloader:
>>>     view0, view1 = views  # each view is a tensor of shape (batch_size, channels, height, width)

__call__(batch: List[Tuple[List[Tensor], int, str]]) → Tuple[List[Tensor], Tensor, List[str]]

Turns a batch of (views, label, filename) tuples into a single (views, labels, filenames) tuple.

Parameters: batch – The input batch as a list of (views, label, filename) tuples, one for each image in the batch. views is a list of N view tensors, each representing a transformed version of the original image. label and filename are the class label and filename for the corresponding image.

Example

>>> batch = [
>>>     ([img_0_view_0, ..., img_0_view_N], label_0, filename_0),   # image 0
>>>     ([img_1_view_0, ..., img_1_view_N], label_1, filename_1),   # image 1
>>>     ...
>>>     ([img_B_view_0, ..., img_B_view_N], label_B, filename_B),  # image B
>>> ]

Returns

A tuple containing –

views: A list of tensors, where each tensor corresponds to one view of every image in the batch. Tensors are concatenated along the batch dimension.
labels: A tensor of shape (batch_size,) with torch.long dtype, containing the labels for all images in the batch.
filenames: A list of strings containing filenames for all images in the batch.

Example

>>> output = (
>>>     [
>>>         Tensor([img_0_view_0, ..., img_B_view_0]),    # view 0
>>>         Tensor([img_0_view_1, ..., img_B_view_1]),    # view 1
>>>         ...
>>>         Tensor([img_0_view_N, ..., img_B_view_N]),    # view N
>>>     ],
>>>     torch.tensor([label_0, ..., label_B], dtype=torch.long),
>>>     [filename_0, ..., filename_B],
>>> )

Notes

If the input batch is empty, a warning is issued, and an empty tuple ([], [], []) is returned.

.ijepa_collate

Note

IJEPAMaskCollator used to live in lightly.data.collate, which has been removed. Use from lightly.data import IJEPAMaskCollator.

class lightly.data.ijepa_collate.IJEPAMaskCollator(input_size: int | tuple[int, int] = (224, 224), patch_size: int = 16, enc_mask_scale: tuple[float, float] = (0.2, 0.8), pred_mask_scale: tuple[float, float] = (0.2, 0.8), aspect_ratio: tuple[float, float] = (0.3, 3.0), nenc: int = 1, npred: int = 2, min_keep: int = 4, allow_overlap: bool = False)

Collator for IJEPA model [0].

Experimental: Support for I-JEPA is experimental, there might be breaking changes in the future.

Code inspired by [1].

[0]: Joint-Embedding Predictive Architecture, 2023, https://arxiv.org/abs/2301.08243
[1]: https://github.com/facebookresearch/ijepa

__call__(batch: list[Any]) → tuple[Any, Any, Any]: Create encoder and predictor masks when collating imgs into a batch # 1. sample enc block (size + location) using seed # 2. sample pred block (size) using seed # 3. sample several enc block locations for each image (w/o seed) # 4. sample several pred block locations for each image (w/o seed) # 5. return enc mask and pred mask