API Reference

Dataset

class dacapo_toolbox.dataset.iterable_dataset(datasets, shapes, weights=None, transforms=None, sampling_strategies=None, trim=None, simple_augment_config=None, deform_augment_config=None, interpolatable=None)

Builds a gunpowder pipeline and wraps it in a torch IterableDataset. See https://pytorch.org/docs/stable/data.html#torch.utils.data.IterableDataset for more info

Parameters:

datasets (dict[str, Array | Graph | Sequence[Array] | Sequence[Graph]]) – A dictionary mapping dataset names to their corresponding data arrays or graphs. Each key can correspond to a single Array/Graph or a sequence of Arrays/Graphs. If a sequence is provided, the datasets will be paired together during sampling based on ordering. I.e. if you have “raw”: [raw1, raw2] and “gt”: [gt1, gt2], then raw1 will always be paired with gt1 and raw2 with gt2. All arrays must have the same relative voxel size (i.e. raw1 with voxel size (4,4,40), raw2 with voxel size (2,2,20), and gt1 with voxel_size (2,2,20) would mean gt2 must have voxel size (1,1,10)) since gt1 has half the voxel size of raw1, gt2 must have half the voxel size of raw2.

shapes (dict[str, Sequence[int]]) – A dictionary mapping dataset names to their corresponding shapes. The output of the dataset will be batches of arrays with these shapes, regardless of the original array shapes and voxel sizes.

weights (Sequence[float] | None) – An optional list of weights defining the sampling rate for entrees in a dataset. I.e. if you have “raw”: [raw1, raw2] and weights=[0.8, 0.2], then 80% of the time raw1 will be sampled and 20% of the time raw2 will be sampled.

transforms (dict[str | tuple[str | tuple[str, ...], str | tuple[str, ...]], Callable] | None) –
An optional dictionary mapping operations to be called on arrays during batch generation. The keys have the form ((“in_key_a”, “in_key_b”, …), (“out_key_a”, “out_key_b”, …)). with support for the following shorthands:

”key” is shorthand for ((“key”), (“key”)) (i.e. in-place transform)

(“key_a”, “key_b”) is shorthand for ((“key_a”,),(“key_b”,)) (i.e. one-to-one transform)
Similarly (“key_a”, (“key_b”, “key_c”)) and ((“key_a”, “key_b”), “key_c”) are supported.

sampling_strategies (MaskedSampling | PointSampling | Sequence[MaskedSampling | PointSampling] | None) –
An optional list of sampling strategies to be applied to each dataset. If a single strategy is provided, it will be applied to all datasets. If None, no sampling strategy will be applied and random sampling will be used. The recommended strategies are as follows:

None - Random sampling, best for densely or almost densely annotated crops.

PointSampling - Best for large sparse annotations, requires a graph of points to sample from. Points should be provided as a networkx graph with node attribute “position” containing the coordinates of the points in world units.

MaskedSampling - Best for small sparse annotations. Will build an integral mask for quickly computing the fraction of the crop that would be contained for any given location, trading memory for speed. If you dataset is large, this may not be feasible, in which case you can use the “reject” strategy, which will randomly sample locations and reject those that do not meet the minimal fraction of mask coverage required. This is extremely inefficient for very sparse annotations.

trim (int | Sequence[int] | None) – We generally guarantee that the center voxel of the requested shape will be contained in the valid region of the dataset, meaning in the worst case scenario, half of the volume could contain padding. If you want to avoid this you can pass in the trim parameter, which will ensure the center voxel is at least trim voxels from the edge of the valid region, limiting the amount of padding that can be introduced.

simple_augment_config (SimpleAugmentConfig | None) – An optional SimpleAugmentConfig object defining the parameters for simple augmentations (mirroring and transposing).

deform_augment_config (DeformAugmentConfig | None) – An optional DeformAugmentConfig object defining the parameters for deform augmentations (scaling, rotation, elastic deformations).

interpolatable (dict[str, bool] | None) – An optional dictionary mapping dataset names to a boolean indicating whether the dataset should be interpolated during augmentations. If not provided, defaults to True for float arrays, uint8, and uint16 arrays (commonly used for raw data), and False for 32/64 bit signed and unsigned ints (commonly used for labels).

Return type:

IterableDataset

Returns:

A torch IterableDataset that can be used with a DataLoader to provide batches of data.

class dacapo_toolbox.dataset.SimpleAugmentConfig(p=0.0, mirror_only=None, transpose_only=None, mirror_probs=None, transpose_probs=None)

The simple augment handles non-interpolating geometric transformations. This includes mirroring and transposing in n-dimensional space. See https://github.com/funkelab/gunpowder/blob/main/gunpowder/nodes/simple_augment.py for more details.

Parameters:

p (float) – Probability of applying the augmentations.

mirror_only (Sequence[int] | None) – List of axes to mirror. If None, all axes may be mirrored.

transpose_only (Sequence[int] | None) – List of axes to transpose. If None, all axes may be transposed.

mirror_probs (Sequence[float] | None) – List of probabilities for each axis in mirror_only. If None, uses equal probability for all axes.

transpose_probs (dict[tuple[int, ...], float] | Sequence[float] | None) – Dictionary mapping tuples of axes to probabilities for transposing. If None, uses equal probability for all axes.

class dacapo_toolbox.dataset.DeformAugmentConfig(p=0.0, control_point_spacing=None, jitter_sigma=None, scale_interval=None, rotate=False, subsample=4, spatial_dims=3, rotation_axes=None)

The deform augment handles interpolating geometric transformations. This includes scaling, rotation, and elastic deformations. See https://github.com/funkelab/gunpowder/blob/main/gunpowder/nodes/deform_augment.py for more details.

Parameters:

p (float) – Probability of applying the augmentations.

control_point_spacing (Sequence[int] | None) – Spacing of the control points for the elastic deformation.

jitter_sigma (Sequence[float] | None) – Standard deviation of the Gaussian noise added to the control points.

scale_interval (tuple[float, float] | None) – Interval for scaling the input data.

rotate (bool) – Whether to apply random rotations.

subsample (int) – Subsampling factor for the control points.

spatial_dims (int) – Number of spatial dimensions.

rotation_axes (Sequence[int] | None) – Axes around which to rotate. If None, rotates around all axes.

class dacapo_toolbox.dataset.MaskedSampling(mask_key, min_masked=1.0, strategy='integral_mask')

Sampling strategy that uses a mask to determine which samples to include.

Parameters:

mask_key (str) – The key of the mask array in the dataset.

min_masked (float) – Minimum fraction of samples that must be masked in to include the sample. If less than this fraction is masked in, the sample is skipped.

strategy (str) – Optional strategy to apply to the mask. by default generates an integral mask for quick sampling at the cost of extra memory usage. If your dataset is large, you may want to use “reject”.

class dacapo_toolbox.dataset.PointSampling(sample_points_key)

Sampling strategy that uses a set of points to determine which samples to include.

Parameters:

sample_points_key (str) – The key of the sample points array in the dataset.

Transforms

Affinities

class dacapo_toolbox.transforms.affs.Affs(neighborhood, dist_func='equality', pad=True, concat_dim=0)

A torch module to compute affinities from a label tensor using a specified neighborhood and distance function.

Parameters:

neighborhood (Sequence[Sequence[int]]) – A sequence of offsets defining the neighborhood for which to compute affinities. Each offset should be a sequence of integers with length equal to the number of spatial dimensions in the input tensor.

dist_func (Union[str, Callable, list[Callable]]) –
A string or callable defining the distance function to use for computing affinities. If a string is provided, it should be one of the following:

”equality”: Affinity is 1 if the labels are equal, 0 otherwise.

”equality-no-bg”: Affinity is 1 if the labels are equal and non-zero, 0 otherwise.

Alternatively, a callable can be provided that takes two tensors as input and returns a tensor of the same shape with affinity values. If a list of callables is provided, each callable will be applied to the corresponding offset in the neighborhood. Finally the callable can be a torch.nn.Module or a list of torch.nn.Modules, allowing for learnable distance functions.

pad (bool) – If True, the input tensor will be padded such that the output tensor has the same shape as the input tensor. If False, the output tensor will be smaller than the input tensor depending on the offsets in the neighborhood.

concat_dim (int) – The dimension along which to concatenate the output affinities. Default is 0, but when working with batched data, you may want to set this to 1 to get an output of shape (batch, channels, …).

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: Tensor

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dacapo_toolbox.transforms.affs.AffsMask(neighborhood, pad=True)

Deprecated, just use Affs with the appropriate distance function to get your mask. It is a dataset specific decision whether you want to mask affinities between fg/bg pairs, bg/bg pairs, both, or neither.

forward(mask)

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: Tensor

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Distances

class dacapo_toolbox.transforms.distances.SignedDistanceTransform(sigma=10.0)

Computes the signed distance transform of a label mask. The output is normalized to the range [-1, 1] using tanh(dist/sigma).

Parameters:

sigma (float) – Controls the steepness of the tanh function. Higher values result in a gentler slope. Default is 10.0.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: Tensor

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dacapo_toolbox.transforms.distances.SDTBoundaryMask

Computes a binary mask of regions where the distance to the nearest boundary is less than the distance to the border of the image. This is useful to avoid training on the ambiguous boundary regions where the true distance to the nearest object is not known.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: Tensor

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

LSDS

class dacapo_toolbox.transforms.lsds.LSD(sigma, voxel_size=None, labels=None, downsample=1)

Compute local shape descriptors (LSD) for a segmentation.

Parameters:

sigma (float | Sequence[float]) – The gaussian kernel sigma to consider for the local shape descriptor.

voxel_size – The voxel size of the segmentation. Defaults to 1 per axis.

labels – Restrict the computation to the given labels. Defaults to all labels inside the ROI of the segmentation.

downsample – Compute the local shape descriptor on a downsampled volume for faster processing. Defaults to 1 (no downsampling).

forward(segmentation)

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: Tensor

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

dacapo_toolbox.transforms.lsds.get_local_shape_descriptors(segmentation, sigma, voxel_size=None, labels=None, downsample=1)

Compute local shape descriptors for the given segmentation.

Parameters:

segmentation (ndarray) – A label array to compute the local shape descriptors for.

sigma (float | Sequence[float]) – The radius to consider for the local shape descriptor.

voxel_size (Sequence[int] | None) – The voxel size of segmentation. Defaults to 1.

labels (Sequence[int] | None) – Restrict the computation to the given labels. Defaults to all labels inside the roi of segmentation.

downsample (int) – Compute the local shape descriptor on a downsampled volume for faster processing. Defaults to 1 (no downsampling).

Weight Balancing

class dacapo_toolbox.transforms.weight_balancing.BalanceLabels(slab=None, num_classes=2, clipmin=0.05, clipmax=0.95)

Computes per-voxel weights to balance the contribution of each class to the loss. The weights are computed per-slab, which can be set with the slab parameter. If slab is None, the entire volume is used to compute the weights. The weights are clipped to the range [clipmin, clipmax] to avoid extreme weights.

Parameters:

slab – The shape of the slab to use for computing the weights. If -1 is provided for a dimension, the entire dimension is used. If None is provided, the entire volume is used. Default is None.

num_classes (int) – The number of classes in the labels. Default is 2.

clipmin (float) – The minimum weight to use. Default is 0.05.

clipmax (float) – The maximum weight to use. Default is 0.95.

forward(labels, mask=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

dacapo_toolbox.transforms.weight_balancing.balance_weights(labels, mask=None, num_classes=2, slab=None, clipmin=0.05, clipmax=0.95)

Visualizations

dacapo_toolbox.vis.preview.gif_2d(arrays, array_types, filename, title, fps=10, overwrite=True, dpi=72, max_size=256, optimize_gif=False, frame_skip=2)

Create a 2D GIF preview of the given arrays.

Parameters:

arrays (dict[str, Array]) – A dictionary of named arrays to visualize. Each array must be 3D, with optional channels.

array_types (dict[str, str]) – A dictionary specifying the type of each array. Supported types are: - “raw”: Grayscale raw data. - “labels”: Integer labels, visualized with a random color map. - “affs”: Affinity graphs, visualized as RGB images. - “pca”: High-dimensional data, visualized using PCA to reduce to 3 channels.

filename (str) – The output filename for the GIF.

title (str) – The title to display on the GIF.

fps (int) – Frames per second for the GIF. Default is 10.

overwrite (bool) – Whether to overwrite the output file if it already exists. Default is True.

dpi (int) – Dots per inch for the output GIF. Default is 72.

max_size (int) – Maximum size (in pixels) for the largest dimension of the images. Default is 256.

optimize_gif (bool) – Whether to optimize the GIF for smaller file size. Default is False.

frame_skip (int) – Number of frames to skip when generating the GIF for faster creation. Default is 2.

dacapo_toolbox.vis.preview.cube(arrays, array_types, filename, title, elev=30, azim=-60, light_azdeg=205, light_altdeg=20, overwrite=True, rcount=128, ccount=128, shade=True, dpi=100)

Preview 3D arrays as cubes with matplotlib. Arrays do not need to be the same size their relative sizes and shifts will be respected.

Parameters:

arrays (dict[str, Array]) – A dictionary of named arrays to visualize. Each array must be 3D, with optional channels.

array_types (dict[str, str]) –
A dictionary specifying the type of each array. Supported types are: - “raw”: Grayscale raw data. - “labels”: Integer labels, visualized with a random color map. - “affs”: Affinity graphs, visualized as RGB images. - “pca”: High-dimensional data, visualized using PCA to reduce to 3

channels.

filename (str) – The output filename for the image.

title (str) – The title to display on the image.

elev (float) – Elevation angle for the 3D view. Default is 30.

azim (float) – Azimuth angle for the 3D view. Default is -60.

light_azdeg (float) – Azimuth angle for the light source. Default is 205.

light_altdeg (float) – Altitude angle for the light source. Default is 20.

overwrite (bool) – Whether to overwrite the output file if it already exists. Default is True.

rcount (int) – Number of rows for the surface plot. Default is 128.

ccount (int) – Number of columns for the surface plot. Default is 128.

shade (bool) – Whether to shade the surface plot. Default is True.

dpi (int) – Dots per inch for the output image. Default is 100.

Sample Data

dacapo_toolbox.sample_datasets.cremi(zarr_path)

Downloads a subset of the CREMI data and returns the raw and label arrays for train and testing.

Parameters:

zarr_path (Path) – The path to the directory where the zarr files will be stored.

Return type:

tuple[Array, Array, Array, Array]

Returns:

A tuple of four arrays (raw_train, labels_train, raw_test, labels_test).