cellmap_flow.finetune.data

What the finetune trainer trains on: random patch pairs read out of a session’s annotation volume, and the loaders that deliver them.

  • dataset: VirtualPatchDataset, the patches. Each is a centre from sampler (the dense, sparse and rehearsal pools), the annotation and raw that reader reads around it, and augment’s flips, rotations, brightness and noise when asked for.

  • loader: dataset_from_manifest and create_dataloader (from a session’s manifest), make_training_loader and rebuild_loader.

Submodules

Classes

VirtualPatchDataset

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

Functions

create_dataloader(→ torch.utils.data.DataLoader)

Build the training DataLoader for a corrections directory.

dataset_from_manifest(...)

Instantiate a VirtualPatchDataset from a manifest dict.

make_training_loader(→ torch.utils.data.DataLoader)

The DataLoader the trainer reads patches through.

rebuild_loader(→ torch.utils.data.DataLoader)

loader again with a different batch size and everything else kept.

Package Contents

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

cellmap_flow.finetune.data.create_dataloader(corrections_zarr_path: str, batch_size: int = 2, augment: bool = True, num_workers: int = 4, patch_voxels: PatchVoxels | None = None) → torch.utils.data.DataLoader

Build the training DataLoader for a corrections directory.

Requires a _virtual_sources.json manifest. Every path that creates a session writes one (volume creation, YAML import, training submit) and restarts backfill one, so a missing manifest means something upstream failed – which is worth an exception rather than training on anything else. The patch geometry comes from the manifest, and the dataset samples randomly, so there is no patch shape or shuffle to pass.

Parameters:

  • corrections_zarr_path – Session corrections directory.

  • batch_size – Clamped down to the dataset size when smaller.

  • augment – Flips, XY rotations, brightness and noise, on top of the patch-centre jitter that is always applied. It overrides the manifest’s stored preference.

  • num_workers – DataLoader workers. Spawned, not forked – tensorstore handles do not survive fork.

  • patch_voxels – (input, output) patch in voxels instead of the manifest’s, for a model that trains on its own (training_patch_voxels).

cellmap_flow.finetune.data.dataset_from_manifest(manifest: dict, corrections_dir: str | None = None, augment: bool | None = None, patch_voxels: PatchVoxels | None = None) → cellmap_flow.finetune.data.dataset.VirtualPatchDataset

Instantiate a VirtualPatchDataset from a manifest dict.

Recognized manifest kinds:

  • volume_zarr_v1 (current): trainer reads the session’s annotation_volume.zarr directly. Field volume_zarr_path.

corrections_dir is where the session’s good regions are looked up; omit it and the dataset simply trains without rehearsal anchors.

patch_voxels, (input, output) in voxels, replaces the manifest’s patch sizes, which are the model’s serving geometry: for a model whose network trains on another (training_patch_voxels). The voxel sizes stay the manifest’s. An epoch then counts as many more patches per annotated chunk as the patch is smaller, so that it still covers what was annotated.

cellmap_flow.finetune.data.make_training_loader(dataset, batch_size: int, num_workers: int) → torch.utils.data.DataLoader

The DataLoader the trainer reads patches through.

Persistent, spawned workers: each keeps its copy of the dataset, and so its RNG, from one epoch to the next. Non-persistent workers are re-spawned every epoch from a fresh pickle of the dataset – RNG unset, same seed – so every epoch would draw the identical patches and augmentations.

cellmap_flow.finetune.data.rebuild_loader(loader: torch.utils.data.DataLoader, batch_size: int) → torch.utils.data.DataLoader

loader again with a different batch size and everything else kept.

The OOM fallback in the trainer used to rebuild its loader with only some of the original arguments; losing persistent_workers made every epoch after an OOM repeat the same patches (see make_training_loader).