cellmap_flow.finetune.correction_dataset

PyTorch Dataset for loading user corrections.

This module provides a Dataset class that loads 3D EM data and correction masks from Zarr files for training LoRA adapters.

Attributes

logger

Classes

CorrectionDataset

PyTorch Dataset for user corrections stored in Zarr format.

Functions

create_dataloader(→ torch.utils.data.DataLoader)

Module Contents

cellmap_flow.finetune.correction_dataset.logger

class cellmap_flow.finetune.correction_dataset.CorrectionDataset(corrections_zarr_path: str, patch_shape: Tuple[int, int, int] | None = None, augment: bool = True, model_name: str | None = None)

PyTorch Dataset for user corrections stored in Zarr format.

Loads raw EM data and corrected masks from corrections.zarr/, with optional 3D augmentation.

Parameters:

corrections_zarr_path – Path to corrections.zarr directory
patch_shape – Shape of patches to extract (Z, Y, X) If None, uses full correction size
augment – Whether to apply 3D augmentation
model_name – If specified, only load corrections for this model

Examples

>>> dataset = CorrectionDataset(
...     "test_corrections.zarr",
...     patch_shape=(64, 64, 64),
...     augment=True
... )
>>> print(f"Dataset size: {len(dataset)}")
>>> raw, target = dataset[0]
>>> print(f"Raw shape: {raw.shape}, Target shape: {target.shape}")

corrections_path

patch_shape = None

augment = True

model_name = None

corrections = []

cellmap_flow.finetune.correction_dataset.create_dataloader(corrections_zarr_path: str, batch_size: int = 2, patch_shape: Tuple[int, int, int] | None = None, augment: bool = True, num_workers: int = 4, shuffle: bool = True, model_name: str | None = None) → torch.utils.data.DataLoader