Datasets#

Load training images and labels, or supply images for prediction. Training datasets augment images and masks together, then generate training targets. See the training guide for a complete loop. From a source checkout, uv sync --locked --extra training installs the required augmentation and HDF5 packages.

Choose the input layout#

Class

Input

Returned image

Other sample keys

TrainDatasetFolder

Parallel image and .mat label folders

image: CHW float32 tensor

Selected mask names; cellpose_inst, etc.

TrainDatasetH5

Configurable HDF5 root arrays

Configured image key: CHW float32 tensor

Selected masks; inst_cellpose, etc.

FolderDatasetInfer

Image folder

image: HWC RGB NumPy array, unless transformed

name: stem without extension

HDF5DatasetInfer

HDF5 /imgs and /fnames

im: CHW tensor retaining stored dtype

file: basename with extension

WSIDatasetInfer

Slide reader and coordinates

image: HWC NumPy array or tensor

name: slide name; coords: x, y, width, height

H is height, W width and C channels. Sample keys differ between datasets, as shown above. Training tensors start on CPU.

Instance labels use 0 for background and a positive ID for each object. Class labels and normalization must match the model checkpoint.

Training from folders#

Image and label files are sorted separately and paired by position. Matching filenames are not checked. Keep both folders aligned, with no extra files. Each .mat file contains selected HW masks such as inst and type.

import albumentations as A

from cellseg_models_pytorch.torch_datasets import TrainDatasetFolder
from cellseg_models_pytorch.transforms.albu_transforms import (
    ApplyEach,
    CellposeTransform,
    MinMaxNormalization,
)

dataset = TrainDatasetFolder(
    "data/train/images",
    "data/train/labels",
    mask_keys=("inst", "type"),
    transforms=A.Compose([MinMaxNormalization(copy=True)]),
    inst_transforms=ApplyEach([CellposeTransform()], as_list=True),
)
sample = dataset[0]
image = sample["image"]
cell_types = sample["type"].long()
flow_target = sample["cellpose_inst"].float()

map_out_keys selects and renames sample entries. Unlisted entries are dropped, so include image if you need it. For example: {"image": "image", "type": "type_mask", "cellpose_inst": "cellpose"}. See transforms for augmentation and target options.

class TrainDatasetFolder(path_img, path_label, mask_keys, transforms, inst_transforms, map_out_keys=None)#

Load aligned RGB images and MATLAB label masks for training.

Image and label files are sorted separately and paired by position, not matched by filename. Keep both folders in the same order with no unrelated files. Spatial augmentation runs jointly before auxiliary targets are generated.

Choose folders, labels and training transforms.

Parameters:

  • path_img (str) – Folder of RGB images supported by the image reader.

  • path_label (str) – Folder of .mat files, one per image in the same sorted order.

  • mask_keys (Tuple[str, ...]) – Label arrays to read: inst, type, cyto_inst, cyto_type or sem. Must include inst. Masks use height, width axes; instance 0 is background and positive labels identify objects. Class meanings come from your data.

  • transforms (A.Compose) – Albumentations Compose receiving image and a list of masks. Include the normalization expected by the model.

  • inst_transforms (ApplyEach) – ApplyEach with as_list=True, deriving targets from the augmented instance masks. Vector targets are converted to CHW tensors.

  • map_out_keys (Dict[str, str]) – Optional dictionary selecting and renaming output entries. Only listed existing keys are retained. To keep the image, include image explicitly. Auxiliary keys are transform_name_instance_key, such as cellpose_inst.

Raises:

  • ModuleNotFoundError – Albumentations is unavailable.

  • ValueError – mask_keys contains unsupported keys or omits inst. Length access also rejects unequal image/label file counts.

__getitem__(ix)#

Read one augmented training sample.

Parameters:

ix (int) – Position in the separately sorted image and label lists.

Returns:

image is CHW float32, selected labels are HW, and generated targets are CHW (for example cellpose_inst). Use map_out_keys to select or rename entries. H is height, W width and C channels.

Label and target dtypes depend on the transforms. Convert class labels to long and flow/distance targets to float before computing losses. A spatial dimension of size one may be removed from labels; use patches with both spatial dimensions greater than one.

Return type:

Dictionary of CPU tensors

__len__()#

Return the paired sample count.

Raises:

ValueError – Image and label folders contain different numbers of files.

Return type:

int

Training from HDF5#

Install tables (PyTables) and Albumentations for this option. An existing HDF5 file must have root-level arrays with a shared first sample axis:

Example node

Shape

Meaning

image

(N, H, W, 3)

RGB images

inst

(N, H, W)

Instance labels

type

(N, H, W)

Cell-type class labels

N is the number of samples. Node names are configurable. The dataset checks their sample counts and reads one sample at a time.

from cellseg_models_pytorch.torch_datasets import TrainDatasetH5

h5_dataset = TrainDatasetH5(
    "data/train.h5",
    img_key="image",
    inst_keys=("inst",),
    mask_keys=("inst", "type"),
    transforms=A.Compose([MinMaxNormalization(copy=True)]),
    inst_transforms=ApplyEach([CellposeTransform()], as_list=True),
)
h5_sample = h5_dataset[0]
flow_target = h5_sample["inst_cellpose"]

This example reuses the imports above. HDF5 target keys put the instance name first (inst_cellpose); folder datasets use cellpose_inst. Labels are int64; images and generated targets are float32. Set normalization through transforms.

class TrainDatasetH5(path, img_key, inst_keys, mask_keys, transforms, inst_transforms)#

Load images and aligned masks from an existing HDF5 file for training.

Requires PyTables and Albumentations. Configured nodes must be arrays with matching first-axis sample counts. Other metadata nodes are ignored.

Choose HDF5 arrays and training transforms.

Parameters:

  • path (str) – Existing HDF5 file readable by PyTables.

  • img_key (str) – Root image node name, containing (N, H, W, C) RGB images.

  • inst_keys (Tuple[str, ...]) – Root instance nodes, each (N, H, W), used to generate targets. Zero is background; positive labels identify individual objects.

  • mask_keys (Tuple[str, ...]) – Root label nodes retained in the returned sample, each (N, H, W). Include an instance node here if you also need its labels.

  • transforms (A.Compose) – Albumentations Compose for joint image/mask augmentation and checkpoint-specific image normalization.

  • inst_transforms (ApplyEach) – ApplyEach with as_list=True, applied after augmentation.

Raises:

  • ModuleNotFoundError – Albumentations or PyTables is unavailable.

  • ValueError – A selected array has no sample axis or has a different count from img_key. Missing files/nodes raise the underlying reader error.

__getitem__(ix)#

Read one aligned, augmented HDF5 training sample.

Parameters:

ix (int) – Index on the arrays’ shared sample axis.

Returns:

CPU tensor dictionary. The image uses img_key, CHW axes and float32 dtype. Labels use mask_keys and int64 dtype; HW labels are squeezed, including any size-one spatial axes. Auxiliary targets use CHW float32 tensors named instance_key_transform_name, such as inst_cellpose. This naming order differs from TrainDatasetFolder. N is sample count, H height, W width and C channels. Values follow the configured transforms.

Return type:

dict[str, Tensor]

__len__()#

Return the shared sample count of the selected arrays.

Return type:

int

Folder inference#

The optional transform receives the image directly: transform(image). Use a function such as minmax_normalize. For Albumentations Compose, call compose(image=image)["image"] inside your transform.

class FolderDatasetInfer(path, pattern='*', transform=None, n_images=None)#

Load RGB images from a folder for inference.

Samples contain image and name. Without a transform, images are uint8 NumPy arrays with height, width, channels axes and no normalization.

Select images and an optional image transform.

Parameters:

  • path (str | Path) – Existing folder containing images.

  • pattern (str) – Glob pattern evaluated inside the folder. Every selected entry must have a lowercase .jpeg, .jpg or .png suffix.

  • transform (Callable[[ndarray], Any] | None) – Callable receiving an image positionally. Its return value becomes the sample’s image. For Albumentations Compose, adapt the call to supply image by keyword and extract its image result.

  • n_images (int) – Optional limit applied after sorting paths, before suffix checks.

Raises:

ValueError – The path is missing/not a directory, or a selected suffix is unsupported.

__getitem__(ix)#

Read one selected image.

Parameters:

ix (int) – Index in the sorted, optionally limited file list.

Returns:

Dictionary with image (RGB HWC NumPy array, or the transform’s return value) and name (filename stem without extension).

Return type:

dict[str, Any]

__len__()#

Return the number of selected image files.

Return type:

int

HDF5 inference#

from cellseg_models_pytorch.torch_datasets.hdf5_dataset_infer import (
    HDF5DatasetInfer,
)

This reader expects /imgs and /fnames, unlike the configurable training reader. It keeps stored pixel values and dtypes. Apply preprocessing before prediction.

class HDF5DatasetInfer(path, n_images=None, **kwargs)#

Load image tensors and filenames from an existing HDF5 file.

This reader expects root nodes imgs and fnames. It does not use the configurable training schema or apply image normalization.

Choose an HDF5 image file and optional sample limit.

Parameters:

  • path (str | Path) – .h5 or .hdf5 file with imgs shaped (N, H, W, C) and UTF-8 byte filenames in fnames, in matching order.

  • n_images (int) – Optional limit on the stored filename list and dataset length.

  • **kwargs – Additional keyword arguments are accepted but unused.

Raises:

  • ImportError – PyTables is unavailable.

  • ValueError – The file suffix is unsupported. Missing files/nodes raise the underlying PyTables error.

__getitem__(ix)#

Read one HDF5 image and its basename.

Parameters:

ix (int) – Sample index.

Returns:

Dictionary with im as a CHW CPU tensor retaining the stored dtype and pixel values, and file as a decoded filename including extension. These keys differ from FolderDatasetInfer’s image/name keys.

Return type:

dict[str, Tensor | str]

__len__()#

Return the number of selected stored filenames.

Return type:

int

Slide-tile inference#

See slide reading for levels and tile coordinates, or the WSI guide for the full workflow. A normalization callback receives image=tile and returns an HWC image. CUDA tiles require cuCIM and a callback that accepts tensors.

class WSIDatasetInfer(reader, coordinates, level=0, transforms=None, *, read_device=None)#

Read slide regions as HWC images for tiled inference.

Unlike the HDF5 inference reader, this dataset keeps channels last. Use the whole-slide guide for the reader and tile-coordinate workflow.

Choose slide regions, read level and optional normalization.

Parameters:

  • reader (SlideReader) – SlideReader used to read regions and supply the slide name.

  • coordinates (Sequence[tuple[int, int, int, int]]) – Indexable sequence of (x, y, width, height) tuples following SlideReader.read_region’s coordinate contract. Pass a TileCoordinates object’s coordinates field for a plain sequence.

  • level (int) – Slide pyramid level passed to read_region.

  • transforms (Callable[[...], dict[str, ndarray | Tensor]] | None) – Callable receiving image=tile and returning a dictionary with an HWC image. Albumentations Compose follows this contract.

  • read_device (str | device | None) – None reads NumPy tiles on CPU. A CUDA device requests GPU tensors through cuCIM; transforms must then accept and return HWC tensors.

__getitem__(index)#

Read and optionally transform one slide region.

Parameters:

index (int) – Position in the supplied coordinate sequence.

Returns:

Dictionary with image (HWC NumPy array or tensor), name (slide name), and coords (NumPy array holding x, y, width, height). Pixel values and image dtype follow the reader and transform; this dataset does not automatically normalize or convert to channels-first tensors.

Return type:

dict[str, ndarray | Tensor | str]

__len__()#

Return the number of requested regions.

Return type:

int