zea.data.dataloader¶

H5 dataloader for loading images from zea datasets.

Example

import zea

loader = zea.Dataloader(
    file_paths="/path/to/dataset",
    key="data/image/values",
    batch_size=16,
    image_range=(-60, 0),
    normalization_range=(0, 1),
    image_size=(256, 256),
    num_threads=16,
)

for batch in loader:
    # batch is a numpy array of shape (batch_size, 256, 256, 1)
    ...

Functions

generate_h5_indices(file_paths, file_shapes, ...)

Generate indices for h5 files.

Classes

`Dataloader`(file_paths, key, *[, batch_size, ...])	High-performance HDF5 dataloader built on Grain.
`H5DataSource`(file_paths[, key, n_frames, ...])	Thread-safe random-access data source for HDF5 files.

class zea.data.dataloader.H5DataSource(file_paths, key='data/image', n_frames=None, frame_index_stride=1, frame_axis=-1, additional_axes_iter=None, sort_files=True, overlapping_blocks=False, limit_n_examples=None, limit_n_frames=None, offset_n_frames=0, return_metadata=None, cache=False, validate=True, revision=None, lazy=True, on_incomplete_blocks='error', on_missing_metadata='error', axis_selections=None, file_filter=None, **kwargs)[source]¶

Bases: object

Thread-safe random-access data source for HDF5 files.

Implements grain.RandomAccessDataSource protocol (__getitem__ and __len__) so it can be plugged directly into a grain.MapDataset pipeline.

Each worker thread gets its own H5FileHandleCache via threading.local() so h5py file handles are never shared across threads.

Parameters:

file_paths (Union[List[str], str]) – Path(s) to HDF5 directory(ies) or file(s).
key (str) – HDF5 dataset key, e.g. "data/image".
n_frames (int | None) – Number of consecutive frames per sample, or None (default) for single frames without a frame axis. See Dataloader.
frame_index_stride (int) – Stride between frames.
frame_axis (int) – Axis the frame block is placed on in the output. Defaults to -1 so frames land in the channel position for image data; see Dataloader. Unused when n_frames is None.
additional_axes_iter (tuple | None) – Extra axes to iterate over.
sort_files (bool) – Sort files numerically.
overlapping_blocks (bool) – Allow overlapping frame blocks.
limit_n_examples (int | None) – Cap the number of examples (dataset length).
limit_n_frames (int | None) – Cap frames loaded per file.
return_metadata (bool | str | Sequence[str] | None) – Return a (sample, metadata) tuple. See Dataloader.
cache (bool) – Cache loaded samples to RAM.
validate (bool) – Validate dataset against the zea format. Default is True.
revision (str | None) – HuggingFace revision (branch, tag, or commit hash) for hf:// paths.
lazy (bool) – Stream hf:// files instead of downloading them. See Dataloader.
on_incomplete_blocks (str) – "error" or "skip" for files too short to fill a block. See Dataloader.
on_missing_metadata (str) – "error" or "skip" for files that cannot supply a requested return_metadata path. See Dataloader.
axis_selections (dict | None) – Map of {axis: indices} pre-filtering non-frame axes, applied to the requested metadata too. See Dataloader.
file_filter (Callable[[File], bool] | dict | None) – Keep only files whose content matches a predicate. See Dataloader for details. Defaults to None (no filtering).

close()[source]¶

Close all file handles across all threads.

Handles reopen lazily on the next read, so the source stays usable afterwards.

zea.data.dataloader.generate_h5_indices(file_paths, file_shapes, n_frames, frame_index_stride, key='data/image', source_frame_axis=0, additional_axes_iter=None, sort_files=True, overlapping_blocks=False, limit_n_frames=None, on_incomplete_blocks='error', axis_selections=None, offset_n_frames=0)[source]¶

Generate indices for h5 files.

Generates a list of indices to extract images from hdf5 files. Length of this list is the length of the extracted dataset.

Parameters:

file_paths (List[str]) – List of file paths.
file_shapes (list) – List of file shapes.
n_frames (int | None) – Number of frames per sample. None selects single frames with an integer index, so the frame axis is dropped from the result.
frame_index_stride (int) – Interval between frames to load.
key (str) – Key of hdf5 dataset to grab data from. Defaults to “data/image”.
source_frame_axis (int | None) – Axis of the file’s arrays that stores frames, or None when the data has no frame axis, in which case every file yields a single sample. Defaults to 0.
additional_axes_iter (Optional[List[int]]) – Additional axes to iterate over in the dataset. Defaults to None.
sort_files (bool) – Sort files by number. Defaults to True.
overlapping_blocks (bool) – Will take n_frames from sequence, then move by 1. Defaults to False.
limit_n_frames (int | None) – Maximum number of frames to load per file, counted from offset_n_frames. Defaults to None (no limit).
on_incomplete_blocks (str) – What to do with files holding too few frames to fill one block of n_frames frames spaced by frame_index_stride: "error" (default) raises and names them, "skip" drops them from the index table.
axis_selections (dict | None) – Map of {axis: indices} applied at HDF5 read time to pre-filter non-frame axes. For example {1: [0, 2, 5]} loads only those indices along axis 1, avoiding reading unused data from disk. Defaults to None.
offset_n_frames (int) – Frame index to start iteration from within each file. Combined with limit_n_frames this selects the half-open range [offset_n_frames, offset_n_frames + limit_n_frames). Defaults to 0.

Returns:

List of tuples with indices to extract images from hdf5 files.: (file_name, key, indices) with indices being a tuple of slices.

Return type:

list

Example

[
    (
        "/folder/path_to_file.hdf5",
        "data/image",
        (slice(0, 2, 1), slice(None, 256, None), slice(None, 256, None)),
    ),
    (
        "/folder/path_to_file.hdf5",
        "data/image",
        (slice(2, 4, 1), slice(None, 256, None), slice(None, 256, None)),
    ),
    ...,
]

With n_frames=None the frame entry is a plain int instead of a slice, so the frame axis never enters the loaded array.