cellmap_flow.finetune.data.dataset

The finetune training dataset: random patch pairs out of a session’s annotation volume.

There is exactly one source of truth per session: an annotation_volume.zarr (sparse, full-dataset extent, OME-NGFF) that holds every annotation – painted scribbles plus any imported YAML crops, all merged at their physical offsets. Patches are read straight out of it: no per-tile materialization, no parallel source list to keep in sync.

A patch is a centre drawn from the pools (sampler.py: dense crops, sparse scribbles, rehearsal regions), the annotation and raw read around it (reader.py), and, when asked for, augmentation (augment.py).

A one-slice patch, a 2D network’s, is cut in the plane its paint was drawn in, x–y, z–y or z–x, where the voxels are cubes (slice_planes): there a side slice looks as an x–y one does. Elsewhere it is always an x–y slice, and paint drawn in a side plane, of which it sees a line, is warned about.

Contract with the trainer:

  • len(self) is patches_per_epoch, the epoch length; it has no fixed relation to the number of populated chunks.

  • An item is (raw, annotation), float32 tensors of shape (1, Z, Y, X), and a third, the anchor mask, once the session has a usable good region (emits_anchor).

  • class_voxels is (background, foreground): how many voxels of the whole volume are annotated as each, which the trainer’s single-class warning goes by rather than by the patches an epoch happens to draw.

  • Loader workers are spawned, and each gets the dataset as a pickle, by its class path. The pools go with it; the arrays are opened in the worker.

Attributes

logger

Classes

VirtualPatchDataset

Yield random raw+annotation patches anchored on annotated voxels in a volume zarr.

Functions

stored_voxel_size(raw_dataset_path, voxel_size, resample) → Optional[tuple]

The voxel size (nm) of the raw level patches are read from at

Module Contents

cellmap_flow.finetune.data.dataset.logger
cellmap_flow.finetune.data.dataset.stored_voxel_size(raw_dataset_path: str, voxel_size, resample: bool) → tuple | None

The voxel size (nm) of the raw level patches are read from at voxel_size, as stored, before any resampling to it; None when its metadata cannot be read here.

The level is the one PatchReader’s open_raw opens: chosen for voxel_size as a resampling read chooses (on_mismatch="resample") or as a plain one does, and kept at its own size. Resampled to cubes, voxels stored three times as deep are still blurred across z, so a side slice through them is not an x–y slice.

class cellmap_flow.finetune.data.dataset.VirtualPatchDataset(volume_zarr_path: str, raw_dataset_path: str, input_size_voxels: Tuple[int, int, int], output_size_voxels: Tuple[int, int, int], input_voxel_size_nm: Tuple[float, float, float], output_voxel_size_nm: Tuple[float, float, float], patches_per_epoch: int | None = None, patches_per_chunk: int = 1, jitter_voxels: Tuple[int, int, int] | None = None, seed: int = 0, input_norm_config: dict | None = None, dense_to_sparse_ratio: float | None = None, good_regions: list | None = None, rehearsal_fraction: float | None = None, anchor_fraction: float | None = None, augment: bool = False, resample: bool = False)

Yield random raw+annotation patches anchored on annotated voxels in a volume zarr.

Parameters:

  • volume_zarr_path – path to the session’s annotation_volume.zarr.

  • raw_dataset_path – path to the raw EM zarr the volume is aligned to.

  • input_size_voxels – shape (Z, Y, X) of the raw patch returned per sample, in voxels at input_voxel_size_nm.

  • output_size_voxels – shape (Z, Y, X) of the annotation patch, in voxels at output_voxel_size_nm.

  • input_voxel_size_nm – voxel size for raw patches (the dataset’s closest scale to the model’s claimed input voxel size).

  • output_voxel_size_nm – voxel size for annotation patches.

  • patches_per_epoch – len(self); controls how many random patches comprise one epoch. None (the default) means “auto: substitute the number of annotated chunks” – every annotated chunk gets ~one patch per epoch on average.

  • patches_per_chunk – how many patches the auto epoch length counts per annotated chunk (default 1). A patch much smaller than the chunks (a model trained on tiles, see loader.dataset_from_manifest’s patch_voxels) takes more to cover them.

  • jitter_voxels – half-range of the random offset applied to the patch center, in annotation voxels, in the patch’s order (for a slice cut across y or x, its first entry is across the slice). Defaults to output_size_voxels // 4.

  • seed – RNG seed; per-worker offset added so multi-worker dataloaders sample distinct streams.

  • input_norm_config – the session’s input normalization, applied to every raw patch (see reader.input_normalizers).

  • dense_to_sparse_ratio – fraction in [0, 1] of patches drawn from the dense pool (FG voxels inside any imported_crops bbox). None (default) means auto: 0.5 if both pools have voxels, else 1.0 (use the non-empty pool exclusively).

  • good_regions – the session’s good regions, {"offset_nm", "shape_nm"} boxes to hold the model to its teacher in.

  • rehearsal_fraction – the share of patches centred on a good region. None (default) means a quarter, when any region is usable.

  • anchor_fraction – the share of the other patches centred at a random point of the volume and held to the teacher there. None (default) means a quarter for a painted session, none otherwise.

  • augment – flips, XY rotations, brightness and noise (augment.py).

  • resample – read the raw resampled to input_voxel_size_nm when it has no level at it (the volume’s resample).

volume_zarr_path
raw_dataset_path
augment = False
seed = 0
not_cubes
slice_planes = None
sampler
reader
augmentation
patches_per_epoch: int
property emits_anchor: bool

Whether __getitem__ yields the 3-tuple (raw, annotation, anchor).

Only once there is something to anchor on: good regions, or random anchors (a painted session’s default). Otherwise the dataset keeps its original 2-tuple contract.

property class_voxels: Tuple[int, int]

the voxels annotated as each (PatchSampler.class_voxels).

Type:

(background, foreground)