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).

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).

  • 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.

Module Contents

cellmap_flow.finetune.data.dataset.logger
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. 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
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.