zea.Dataloader¶

class zea.Dataloader(file_paths, key, *, batch_size=16, n_frames=None, shuffle=True, return_metadata=None, seed=None, limit_n_examples=None, limit_n_frames=None, offset_n_frames=0, 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, dtype=None, dataset_repetitions=None, cache=False, additional_axes_iter=None, sort_files=True, overlapping_blocks=False, augmentation=None, on_incomplete_blocks='error', on_missing_metadata='error', frame_index_stride=1, frame_axis=-1, validate=True, revision=None, lazy=True, shard_index=None, num_shards=1, num_threads=16, prefetch_buffer_size=500, reshuffle_each_epoch=True, convert_to_tensor=True, axis_selections=None, file_filter=None, **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):
- offset_n_frames / axis_selections (applied at HDF5 read time)
- limit_n_frames
- limit_n_examples
- shuffle
- shard
- add channel dim
- clip image range
- assert image range
- resize
- repeat
- batch
- cast to dtype (if specified)
- 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.
batch_size (int | None) – Batch size. Set to None to disable batching. Default is 16. Stacking two or more samples (incl. metadata) requires them to have the same shape. This is checked when the loader is built. Note that image_size can resolve differing sample shapes; for the rest, use batch_size=None (or 1, which stacks nothing) and batch it yourself.
n_frames (int | None) – Number of consecutive frames per sample, placed on frame_axis. Default is None, which loads single frames without a frame axis, so a sample keeps the file’s own layout for one frame. Set an int to group consecutive frames into blocks – including n_frames=1, which gives a length-1 frame axis. Frames are read from whichever axis the zea file spec names n_frames for key.
shuffle (bool) – Shuffle dataset each epoch. Default is True.
return_metadata (bool | str | Sequence[str] | None) –
Return a (sample, metadata) tuple instead of a bare sample. False (default) returns arrays only. True returns just the file identity. An iterable of dotted paths additionally loads those fields from the file, e.g. ["scan.sampling_frequency", "metadata.subject"]; a path pointing at a group loads everything below it. Paths use the same syntax as file_filter. The returned dict mirrors FileSpec, with the loader’s own provenance under a "file" key:
```
{
    "scan": {"sampling_frequency": 40e6},
    "metadata": {"subject": {"age": 61}},
    "file": {"fullpath": ..., "filename": ..., "indices": ...},
}
```
Fields whose leading dimension is n_frames in the spec are sliced to the sample’s frames so they stay aligned with the returned images, and fields sharing a dimension with an axis_selections entry are narrowed the same way. During construction, the loader checks that all files can supply the requested paths and that they have the same shapes, raising KeyError or ValueError naming the offending files. They can be dropped with on_missing_metadata="skip" or file_filter. Use batch_size=None for metadata that genuinely varies in shape between files.
seed (int | None) – Random seed used for dataloader (e.g. shuffling). Default is None. If None a random seed is generated.
limit_n_examples (int | None) – Cap the total number of examples (== item before batching) the loader yields, across all files (useful for debugging). Default is None (no limit). Note that this happens before shuffle.
limit_n_frames (int | None) – Maximum number of frames to load per file, counted from offset_n_frames. Default is None (no limit).
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). Default is 0.
drop_remainder (bool) – Drop the final incomplete batch. Default is False.
image_size (tuple | None) – Target (height, width). Default is None (no resizing). Setting it is what lets files of differing image size be batched together, since they arrive at the batch op already sharing a shape.
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.
dtype (str | dtype | None) – Cast samples to this dtype (e.g. "float32", np.float16) after batching and before normalization. Must be floating point whenever normalization_range is set. Default is None, which keeps the dtype the files hold – except that files holding integers are promoted to float32, since normalizing has no integer-valued result.
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, on top of the frame axis. Each becomes an integer index, so those axes are dropped from the sample. 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 unless n_frames > 1. Default is False.
on_incomplete_blocks (str) – What to do with files holding too few frames to fill one block of n_frames (spaced by frame_index_stride). "error" (default) refuses to build the loader and names the offending files; "skip" drops them from the dataset.
on_missing_metadata (str) – What to do with files that cannot supply a path requested through return_metadata. "error" (default) refuses to build the loader and names the offending files; "skip" drops them from the dataset. Default is "error".
augmentation (Callable | None) – Callable applied to each batch after normalization. Default is None.
frame_index_stride (int) – Step between selected frames, for example, 2 takes every other frame, 3 every third. Samples still follow one another without gaps – with n_frames=2, frame_index_stride=2 a file yields frames (0, 2), then (3, 5). Default is 1.
frame_axis (int) – The frames are put in this axis. Only applies when n_frames is set. Default is -1: an image batch comes out as (batch, height, width, n_frames), the channels-last layout Resizer and keras expect. That is why resizing without explicit resize_axes requires frame_axis=-1. For raw_data it makes sense to set frame_axis=0 to keep frames in front.
validate (bool) – Validate discovered files against the zea format, raising if any file is not a valid zea file. Default is True. The verdict is cached, so only the first run over a given dataset opens every file.
revision (str | None) – HuggingFace revision (branch, tag, or commit hash) for hf:// paths. Defaults to None (uses the default branch, typically "main").
lazy (bool) – Stream hf:// files over HTTP instead of downloading them up front. Default is True: a read fetches only the chunks it touches, so reading a slice of a few large files costs a fraction of them, and the bytes land in the on-disk chunk cache (zea.data.chunk_cache) for the next read. Pass False to download every file in full before the pipeline is built, which is the better trade for training: shuffled access over a dataset larger than that cache re-fetches it every epoch, and a bulk download is both faster per byte and resumable. Ignored for local paths.
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) – How many elements the Grain buffer holds, per Python process (not per thread). An element here is whatever this loader yields, so with a batch_size it counts batches: prefetch_buffer_size=500, batch_size=16 keeps up to 8000 examples in RAM. Useful when reading from a distributed file system. Default is 500. Set to 0 to disable prefetching and read sequentially instead.
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.
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 chunks from disk. This can save time and memory. return_metadata fields carrying the selected dimension are cut to match, wherever that dimension sits in their own layout: selecting transmits on data/raw_data also selects them in scan.t0_delays. Default is None.
file_filter (Callable[[File], bool] | dict | None) – Keep only files whose content matches a predicate, discarding the rest before any frames are indexed. Either a callable File -> bool (a file is kept when it returns True), or a declarative dotted-path dict mapping a path on the File to a condition: the EXISTS() helper (field must be present), a plain value (equality), or a callable on the resolved value. All dict entries are ANDed. Files whose predicate raises (e.g. they have no metadata group) are excluded. Default is None (no filtering).

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)

Filtering examples:

from zea import Dataloader, EXISTS

# callable: keep only files that record a subject fat percentage
loader = Dataloader(
    file_paths="filter-demo-dataset",
    key="data/image/values",
    file_filter=lambda f: f.metadata.subject is not None
    and f.metadata.subject.fat_percentage is not None,
)

# metadata: load selected fields alongside each sample
loader = Dataloader(
    file_paths="filter-demo-dataset",
    key="data/image/values",
    batch_size=None,
    return_metadata=["scan.center_frequency", "metadata.subject"],
)
sample, meta = next(iter(loader))
assert meta["scan"]["center_frequency"] in (5e6, 9e6)
assert meta["metadata"]["subject"]["sex"] in ("f", "m")
assert meta["file"]["filename"] in ("a", "b")

# dict: presence + equality + a value-level predicate (all ANDed)
loader = Dataloader(
    file_paths="filter-demo-dataset",
    key="data/image/values",
    file_filter={
        "metadata.subject.fat_percentage": EXISTS,
        "metadata.subject.sex": "f",
        "scan.center_frequency": lambda v: 4e6 <= v <= 6e6,
    },
)

close()[source]¶: Release file handles.

property dataset¶: The underlying grain.MapDataset.

property sample_shapes: dict[tuple, list[str]]¶

Each shape this loader yields, mapped to a few files that produce it.

One entry for a well-formed dataset; more than one means the files disagree and only batch_size=None can iterate them.

property shape¶

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

With drop_remainder=False the final batch is shorter than the rest whenever the sample count is not a multiple of batch_size; the batch axis is reported as None in that case, since no single size describes every batch.

Raises:: ValueError – If the loader yields more than one shape, which only an unbatched loader can do – there is no single shape to report. Inspect sample_shapes for the shapes it does yield.

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