zea.Dataloader

class zea.Dataloader(file_paths, key='data/image', batch_size=16, n_frames=1, shuffle=True, return_filename=False, seed=None, limit_n_samples=None, limit_n_frames=None, drop_remainder=False, image_size=None, resize_type=None, resize_axes=None, resize_kwargs=None, image_range=None, normalization_range=None, clip_image_range=False, assert_image_range=True, dataset_repetitions=None, cache=False, additional_axes_iter=None, sort_files=True, overlapping_blocks=False, augmentation=None, initial_frame_axis=0, insert_frame_axis=True, frame_index_stride=1, frame_axis=-1, validate=True, revision=None, prefetch=True, shard_index=None, num_shards=1, num_threads=16, prefetch_buffer_size=500, reshuffle_each_epoch=True, convert_to_tensor=True, **kwargs)[source]

Bases: object

High-performance HDF5 dataloader built on Grain.

grain threads (N) → h5py (thread-local handles) → numpy -> cpu tensor → user

The entire pipeline runs using numpy, and the resizing is done on the selected backend, all on cpu.

Does the following in order to load a dataset:

  • Find all .hdf5 files in the director(ies)

  • Load the data from each file using the specified key

  • Apply the following transformations in order (if specified):

    • shuffle

    • shard

    • add channel dim

    • clip image range

    • assert image range

    • resize

    • repeat

    • batch

    • cast to float32

    • normalize

    • augmentation

    • convert_to_tensor

Parameters:

  • file_paths (Union[List[str], str]) – Path(s) to directory(ies) and/or HDF5 file(s).

  • key (str) – HDF5 dataset key. Default is "data/image".

  • batch_size (int | None) – Batch size. Set to None to disable batching. Default is 16.

  • n_frames (int) – Number of consecutive frames per sample. Default is 1. When n_frames > 1, frames are grouped into blocks.

  • shuffle (bool) – Shuffle dataset each epoch. Default is True.

  • return_filename (bool) – Return filename metadata together with each sample. Default is False.

  • seed (int | None) – Random seed used for dataloader (e.g. shuffling). Default is None. If None a random seed is generated.

  • limit_n_samples (int | None) – Limit total number of samples (useful for debugging). Default is None (no limit).

  • limit_n_frames (int | None) – Limit frames loaded per file to the first N frames. Default is None (no limit).

  • drop_remainder (bool) – Drop the final incomplete batch. Default is False.

  • image_size (tuple | None) – Target (height, width). Default is None (no resizing).

  • resize_type (str | None) – Resize strategy. One of "resize", "center_crop", "random_crop" or "crop_or_pad". Default is None, which resolves to "resize" when image_size is set.

  • resize_axes (tuple | None) – Axes to resize along, must have length 2 (height, width). Only needed when data has more than (h, w, c) dimensions. Axes are interpreted after frame-axis insertion/reordering. Default is None.

  • resize_kwargs (dict | None) – Extra keyword arguments passed to Resizer. Default is None.

  • image_range (tuple | None) – Source value range of images, e.g. (-60, 0). Used for clipping/asserting/normalization. Default is None.

  • normalization_range (tuple | None) – Target value range, e.g. (0, 1). If set, image_range must also be set. Default is None.

  • clip_image_range (bool) – Clip values to image_range before normalization. Default is False.

  • assert_image_range (bool) – Assert values stay within image_range. Default is True.

  • dataset_repetitions (int | None) – Repeat dataset this many times. Repetition happens after sharding. Default is None (no repetition).

  • cache (bool) – Cache loaded samples in RAM. Default is False. Note that with overlapping_blocks=True, the same frame can be part of multiple samples, so caching will consume more memory.

  • additional_axes_iter (tuple | None) – Additional axes to iterate over in addition to initial_frame_axis. Default is None.

  • sort_files (bool) – Sort files numerically before indexing. Default is True.

  • overlapping_blocks (bool) – If True, frame blocks overlap by n_frames - 1. Has no effect when n_frames == 1. Default is False.

  • augmentation (callable) – Callable applied to each batch after normalization. Default is None.

  • initial_frame_axis (int) – Axis in file data that represents frames. Default is 0.

  • insert_frame_axis (bool) – If True, keep per-frame samples and move/insert the frame dimension at frame_axis. If False, loaded frames are concatenated along frame_axis. Default is True.

  • frame_index_stride (int) – Step between selected frames in a block. Default is 1.

  • frame_axis (int) – Axis along which frames are stacked/placed in output. Default is -1.

  • validate (bool) – Validate discovered files against the zea format. Default is True.

  • revision (str | None) – HuggingFace revision (branch, tag, or commit hash) for hf:// paths. Defaults to None (uses the default branch, typically "main").

  • prefetch (bool) – Enable Grain prefetching for iteration. Default is True.

  • shard_index (int | None) – Shard index to select when num_shards > 1. Must satisfy 0 <= shard_index < num_shards.

  • num_shards (int) – Total number of shards for distributed loading. Sharding happens before downstream transforms. Default is 1.

  • num_threads (int) – Number of Grain read threads (0 means main thread only). Default is 16.

  • prefetch_buffer_size (int) – Size of the Grain buffer for reading elements per Python process (not per thread). Useful when reading from a distributed file system. Default is 500.

  • reshuffle_each_epoch (bool) – Whether to reshuffle the dataset after each epoch. Default is True. For evaluation it might be useful to set this to False. Or when you want to use a persistent iterator between epochs, using dataset_repetitions to specify the number of epochs.

  • convert_to_tensor (bool) – Whether to convert the data to a tensor (on cpu). Default is True.

Example

loader = Dataloader(
    file_paths="/data/camus",
    key="data/image/values",
    batch_size=32,
    image_range=(-60, 0),
    normalization_range=(0, 1),
    image_size=(256, 256),
)
for batch in loader:
    ...  # batch.shape == (32, 256, 256, 1)
close()[source]

Release file handles.

property dataset

The underlying grain.MapDataset.

property shape

Output shape of one batch (or sample if unbatched).

shuffle(seed=None)[source]

(Re-)shuffle the dataset. Rebuilds the pipeline with a fresh seed.

summary()[source]

Print dataset statistics and per-directory breakdown.

to_iter_dataset()[source]

Convert to a grain.IterDataset with prefetching.

This is called automatically when you iterate, but you can call it explicitly if you want to hold onto the IterDataset object.

Return type:

IterDataset