miao

Scalable PyTorch data loaders for OME-NGFF zarr datasets, powered by TensorStore.

pip install miao-io

1. Write a config

# config.yaml
volumes:
  - name: "raw"
    path: "/data/sample_A.zarr"
    image_key: "raw"
    zarr_version: "zarr2"      # or "zarr3" (default: "zarr2")
    label_key: "labels/seg"    # optional
    weight: 0.5                # optional

  - name: "membrane"
    path: "/data/sample_B.zarr"
    image_key: "predictions"
    zarr_version: "zarr3"
    weight: 0.2

  - name: "expanded"
    path: "/data/expanded.zarr"
    image_key: "raw"
    exp_factor: 4.0            # optional, default: 1.0 — see Effective voxel sizes
    weight: 0.3

resolutions: [[8, 8, 8], [16, 16, 16], [32, 32, 32]]  # effective output voxel size per scale
output_axes: "lcxyz"            # layer, channels, X, Y, Z. Shuffle as you please!!!
patch_size: [64, 64, 64]
samples_per_epoch: 1000
cache_bytes: 1073741824         # 1 GB tensorstore cache

2. Build a dataset

from torch.utils.data import DataLoader
from miao import VolumeDataset, load_config

dataset = VolumeDataset(load_config("config.yaml"))
loader = DataLoader(dataset, batch_size=4, shuffle=True, num_workers=8,
                    pin_memory=True, persistent_workers=True)

for batch in loader:
    img, label = batch["img"], batch["label"]

Runnable notebooks live in examples/.

Spatial axes follow output_axes throughout. pixel_size sits at the top level rather than inside meta so default collate stacks it cleanly, ready to feed a scale-conditioned model. It reports the resolutions actually used for that sample, which under random resolution sampling differ from sample to sample.

Scales are addressed by desired output resolution — physical voxel size per axis, in the zarr's OME coordinateTransformations unit (e.g. nanometers) — never by pyramid level index. One resolutions list therefore serves volumes with wildly different pyramid layouts, and any volume can override it with its own.

Every voxel size and resolution in this README is an effective voxel size: the size of one voxel in the specimen, not in the microscope. For a volume with exp_factor: F, the effective voxel size is the stored coordinateTransformations scale divided by F — a 4× expanded sample imaged at 200 nm holds 50 nm of biology per voxel. exp_factor defaults to 1.0, so for unexpanded data the stored and effective sizes are identical.

Level selection, read extents, pixel_size, and meta["resolutions"] are all in effective units, which is what lets one resolutions list mix expanded and unexpanded volumes correctly. bounding_box and meta["coordinate"] are level-0 voxel indices, not physical sizes, so exp_factor does not affect them.

Per sample, miao:

1. draws a volume by weight, then a random coordinate in its finest-scale (level-0) space;
2. for each requested resolution, picks the coarsest pyramid level whose effective voxel size is still ≤ the target on every axis, preferring downsampling — or the finest stored level, upsampled, if the target is finer than anything on disk;
3. reads ceil(patch_size × target_resolution / effective_level_voxel_size) voxels centered on that coordinate and resamples them to patch_size — trilinear for images, nearest for labels.

Every crop thus holds the same voxel count while covering a wider physical extent at coarser scales.

To train resolution-agnostic models, let each sample draw its own scales. Set resolution_sampling instead of resolutions, globally or per volume:

resolution_sampling:
  strategy: log_uniform     # only option for now (pluggable; gaussian etc. can be added)
  ranges: [[[8], [64]]]     # one or more [min, max] ranges
  n_scales: 3               # scales per range — scalar (all ranges) or list (one per range)
  sort: true                # sort drawn scales fine -> coarse

The shape of each bound decides isotropy:

• [[v], [w]] is isotropic — one value per scale, broadcast to all axes. [[[1], [4]]] gives cubic voxels between 1 and 4.
• [[a, b, c], [d, e, f]] draws each axis independently, so [[[1, 1, 1], [4, 4, 4]]] may yield [1.8, 3.2, 2.5] — not the same distribution as [[[1], [4]]].

Ranges stack: [[[1], [2]], [[4], [8]]] with n_scales: [1, 1] gives two scales, the l dimension being the sum across ranges. Draws are log-uniform and redrawn every __getitem__.

Exactly one of resolutions / resolution_sampling may be set. resolution_sampling rules out sampling: "sequential", and sample_windows demands isotropic ranges.

Random sampling — the default, and the mode described above — draws samples_per_epoch patches per epoch. Two settings change where patches land:

Walk the whole volume on a deterministic grid instead.

sampling: "sequential"
overlap: 16              # or per-axis list, e.g. [16, 16, 8]

dataset = VolumeDataset(config)                            # len() = grid positions, all volumes
loader = DataLoader(dataset, batch_size=4, shuffle=False)   # shuffle=False required

for batch in loader:
    # e.g. (2, 0, 3) — grid position per axis; use it to stitch predictions into a full volume
    grid_index = batch["meta"]["grid_index"]

The grid tiles the volume at the first scale's target resolution, striding one output patch minus overlap worth of physical extent. Coverage is gap-free even on anisotropic source data. samples_per_epoch and weight are ignored here, and volumes are exhausted in order.

By default all scales share one center, nesting finer patches inside coarser ones. With sample_windows: true each coarser level instead drops its patch origin uniformly at random among the positions that still cover the previous level's patch.

sample_windows: true

Those sampled origins stay strictly inside any per-volume bounding_box. resolutions must run finest to coarsest (non-decreasing voxel size per axis, e.g. [[8,8,8], [16,16,16]]); putting a coarser resolution first raises.

Set aug_rot: true on a volume to reorient every sample at random:

volumes:
  - name: "raw"
    path: "/data/sample_A.zarr"
    image_key: "raw"
    label_key: "labels/seg"
    aug_rot: true            # optional, default: false

Each __getitem__ draws one transform uniformly from the 48 symmetries of the cube (3! axis permutations × 2³ per-axis flips) — the largest augmentation set reachable without interpolation, since every such transform maps voxels exactly onto voxels. The draw applies to image and labels together, and to every scale, so alignment and nesting survive.

Isotropic output is required, since permuting axes only means something when voxels are cubes. Two checks enforce it: patch_size must be equal across x, y, z (at config load), and the resolution being read must be equal on all axes (per sample, because resolution_sampling can draw a fresh one each call — use isotropic single-value bounds). Anisotropic data on disk is fine — e.g. coarse z upsampled to match — as long as the requested output is isotropic.

Note: aug_rot reorients img and label only. bbox and meta["coordinate"] still describe the original read location; pixel_size is isotropic by construction, so unaffected.

Input axes come from OME-NGFF metadata (multiscales.axes) and never need specifying; channel dimensions are picked up automatically when present.

2D models are often strong baselines in 3D domains. Build one from three degenerate VolumeDatasets — patch_size (1,P,P), (P,1,P), (P,P,1) — glued with ConcatDataset.

from torch.utils.data import ConcatDataset, DataLoader, default_collate
from miao import MiaoConfig, VolumeDataset, load_config

base = load_config("config.yaml")
P = 128
spatial_order = "".join(c for c in base.output_axes if c in "xyzt")   # "xyz"

def plane_config(cfg: MiaoConfig, thin_axis: str, size: int) -> MiaoConfig:
    """Config for a one-voxel-thick patch, i.e. a plane normal to `thin_axis`."""
    data = cfg.model_dump()
    data["patch_size"] = [1 if c == thin_axis else size for c in spatial_order]
    data["samples_per_epoch"] = cfg.samples_per_epoch // len(spatial_order)
    return MiaoConfig(**data)

dataset = ConcatDataset(
    [VolumeDataset(plane_config(base, ax, P)) for ax in spatial_order]
)

def make_plane_collate(output_axes: str):
    """Collate 1-voxel-thick samples by squeezing away the singleton spatial axis."""
    img_dim = {c: i for i, c in enumerate(output_axes) if c in "xyzt"}
    label_axes = output_axes.replace("c", "")     # labels carry no channel axis

    def collate(samples):
        squeezed, planes = [], []
        for sample in samples:
            img = sample["img"]
            thin = [c for c, d in img_dim.items() if img.shape[d] == 1]
            if len(thin) != 1:
                raise ValueError(f"expected one singleton spatial axis, found {thin}")
            axis = thin[0]
            sample = dict(sample)
            sample["img"] = img.squeeze(img_dim[axis])
            if sample["label"].numel():           # empty when a volume has no labels
                sample["label"] = sample["label"].squeeze(label_axes.index(axis))
            squeezed.append(sample)
            planes.append("".join(c for c in img_dim if c != axis))
        batch = default_collate(squeezed)
        batch["plane"] = planes                   # e.g. "yz" — one entry per sample
        return batch

    return collate

loader = DataLoader(dataset, batch_size=8, shuffle=True,
                    collate_fn=make_plane_collate(base.output_axes))

for batch in loader:
    img = batch["img"]                # (B, L, C, P, P)
    label = batch["label"]            # (B, L, P, P)
    bbox = batch["bbox"]              # (B, L, 2, Nd_spatial) — still 3D
    plane = batch["plane"]            # "yz" / "xz" / "xy", one per sample

• Python >= 3.10
• PyTorch >= 2.0
• TensorStore >= 0.1.60
• Zarr datasets following the OME-NGFF spec (v2 or v3)