careamics 0.0.2__py3-none-any.whl → 0.0.3__py3-none-any.whl

careamics/config/likelihood_model.py

@@ -0,0 +1,43 @@
+"""Likelihood model."""
+
+from typing import Literal, Optional, Union
+
+import torch
+from pydantic import BaseModel, ConfigDict
+
+from careamics.models.lvae.noise_models import (
+    GaussianMixtureNoiseModel,
+    MultiChannelNoiseModel,
+)
+
+NoiseModel = Union[GaussianMixtureNoiseModel, MultiChannelNoiseModel]
+
+
+class GaussianLikelihoodConfig(BaseModel):
+    """Gaussian likelihood configuration."""
+
+    model_config = ConfigDict(validate_assignment=True)
+
+    predict_logvar: Optional[Literal["pixelwise"]] = None
+    """If `pixelwise`, log-variance is computed for each pixel, else log-variance
+    is not computed."""
+
+    logvar_lowerbound: Union[float, None] = None
+    """The lowerbound value for log-variance."""
+
+
+class NMLikelihoodConfig(BaseModel):
+    """Noise model likelihood configuration."""
+
+    model_config = ConfigDict(validate_assignment=True, arbitrary_types_allowed=True)
+
+    data_mean: Union[torch.Tensor] = torch.zeros(1)
+    """The mean of the data, used to unnormalize data for noise model evaluation.
+    Shape is (target_ch,) (or (1, target_ch, [1], 1, 1))."""
+
+    data_std: Union[torch.Tensor] = torch.ones(1)
+    """The standard deviation of the data, used to unnormalize data for noise
+    model evaluation. Shape is (target_ch,) (or (1, target_ch, [1], 1, 1))."""
+
+    noise_model: Union[NoiseModel, None] = None
+    """The noise model instance used to compute the likelihood."""

careamics/config/nm_model.py

@@ -0,0 +1,101 @@
+"""Noise models config."""
+
+from pathlib import Path
+from typing import Literal, Optional, Union
+
+import numpy as np
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+from typing_extensions import Self
+
+# TODO: add histogram-based noise model
+
+
+class GaussianMixtureNMConfig(BaseModel):
+    """Gaussian mixture noise model."""
+
+    model_config = ConfigDict(
+        protected_namespaces=(),
+        validate_assignment=True,
+        arbitrary_types_allowed=True,
+        extra="allow",
+    )
+    # model type
+    model_type: Literal["GaussianMixtureNoiseModel"]
+
+    path: Optional[Union[Path, str]] = None
+    """Path to the directory where the trained noise model (*.npz) is saved in the
+    `train` method."""
+
+    signal: Optional[Union[str, Path, np.ndarray]] = None
+    """Path to the file containing signal or respective numpy array."""
+
+    observation: Optional[Union[str, Path, np.ndarray]] = None
+    """Path to the file containing observation or respective numpy array."""
+
+    weight: Optional[np.ndarray] = None
+    """A [3*n_gaussian, n_coeff] sized array containing the values of the weights
+    describing the GMM noise model, with each row corresponding to one
+    parameter of each gaussian, namely [mean, standard deviation and weight].
+    Specifically, rows are organized as follows:
+    - first n_gaussian rows correspond to the means
+    - next n_gaussian rows correspond to the weights
+    - last n_gaussian rows correspond to the standard deviations
+    If `weight=None`, the weight array is initialized using the `min_signal`
+    and `max_signal` parameters."""
+
+    n_gaussian: int = Field(default=1, ge=1)
+    """Number of gaussians used for the GMM."""
+
+    n_coeff: int = Field(default=2, ge=2)
+    """Number of coefficients to describe the functional relationship between gaussian
+    parameters and the signal. 2 implies a linear relationship, 3 implies a quadratic
+    relationship and so on."""
+
+    min_signal: float = Field(default=0.0, ge=0.0)
+    """Minimum signal intensity expected in the image."""
+
+    max_signal: float = Field(default=1.0, ge=0.0)
+    """Maximum signal intensity expected in the image."""
+
+    min_sigma: float = Field(default=200.0, ge=0.0)  # TODO took from nb in pn2v
+    """Minimum value of `standard deviation` allowed in the GMM.
+    All values of `standard deviation` below this are clamped to this value."""
+
+    tol: float = Field(default=1e-10)
+    """Tolerance used in the computation of the noise model likelihood."""
+
+    @model_validator(mode="after")
+    def validate_path_to_pretrained_vs_training_data(self: Self) -> Self:
+        """Validate paths provided in the config.
+
+        Returns
+        -------
+        Self
+            Returns itself.
+        """
+        if self.path and (self.signal is not None or self.observation is not None):
+            raise ValueError(
+                "Either only 'path' to pre-trained noise model should be"
+                "provided or only signal and observation in form of paths"
+                "or numpy arrays."
+            )
+        if not self.path and (self.signal is None or self.observation is None):
+            raise ValueError(
+                "Either only 'path' to pre-trained noise model should be"
+                "provided or only signal and observation in form of paths"
+                "or numpy arrays."
+            )
+        return self
+
+
+# The noise model is given by a set of GMMs, one for each target
+# e.g., 2 target channels, 2 noise models
+class MultiChannelNMConfig(BaseModel):
+    """Noise Model config aggregating noise models for single output channels."""
+
+    # TODO: check that this model config is OK
+    model_config = ConfigDict(
+        validate_assignment=True, arbitrary_types_allowed=True, extra="allow"
+    )
+    noise_models: list[GaussianMixtureNMConfig]
+    """List of noise models, one for each target channel."""

careamics/config/support/supported_activations.py

@@ -24,3 +24,4 @@ class SupportedActivation(str, BaseEnum):
     TANH = "Tanh"
     RELU = "ReLU"
     LEAKYRELU = "LeakyReLU"
+    ELU = "ELU"

careamics/config/support/supported_algorithms.py

@@ -6,15 +6,28 @@ from careamics.utils import BaseEnum
 
 
 class SupportedAlgorithm(str, BaseEnum):
-    """Algorithms available in CAREamics.
-
-    # TODO
-    """
+    """Algorithms available in CAREamics."""
 
     N2V = "n2v"
+    """Noise2Void algorithm, a self-supervised approach based on blind denoising."""
+
     CARE = "care"
+    """Content-aware image restoration, a supervised algorithm used for a variety
+    of tasks."""
+
     N2N = "n2n"
+    """Noise2Noise algorithm, a self-supervised denoising scheme based on comparing
+    noisy images of the same sample."""
+
+    MUSPLIT = "musplit"
+    """An image splitting approach based on ladder VAE architectures."""
+
+    DENOISPLIT = "denoisplit"
+    """An image splitting and denoising approach based on ladder VAE architectures."""
+
     CUSTOM = "custom"
+    """Custom algorithm, used for cases where a custom architecture is provided."""
+
     # PN2V = "pn2v"
     # HDN = "hdn"
     # SEG = "segmentation"

careamics/config/support/supported_architectures.py

@@ -4,17 +4,14 @@ from careamics.utils import BaseEnum
 
 
 class SupportedArchitecture(str, BaseEnum):
-    """Supported architectures.
+    """Supported architectures."""
 
-    # TODO add details, in particular where to find the API for the models
+    UNET = "UNet"
+    """UNet architecture used with N2V, CARE and Noise2Noise."""
 
-    - UNet: classical UNet compatible with N2V2
-    - VAE: variational Autoencoder
-    - Custom: custom model registered with `@register_model` decorator
-    """
+    LVAE = "LVAE"
+    """Ladder Variational Autoencoder used for muSplit and denoiSplit."""
 
-    UNET = "UNet"
-    VAE = "VAE"
-    CUSTOM = (
-        "Custom"  # TODO all the others tags are small letters, except the architect
-    )
+    CUSTOM = "custom"
+    """Keyword used for custom architectures provided by users and only compatible
+    with `FCNAlgorithmConfig` configuration."""

careamics/config/support/supported_losses.py

@@ -22,6 +22,8 @@ class SupportedLoss(str, BaseEnum):
     N2V = "n2v"
     # PN2V = "pn2v"
     # HDN = "hdn"
+    MUSPLIT = "musplit"
+    DENOISPLIT = "denoisplit"
+    DENOISPLIT_MUSPLIT = "denoisplit_musplit"
     # CE = "ce"
     # DICE = "dice"
-    # CUSTOM = "custom" # TODO create mechanism for that

careamics/config/transformations/n2v_manipulate_model.py

@@ -33,7 +33,7 @@ class N2VManipulateModel(TransformModel):
 
     name: Literal["N2VManipulate"] = "N2VManipulate"
     roi_size: int = Field(default=11, ge=3, le=21)
-    masked_pixel_percentage: float = Field(default=0.2, ge=0.05, le=1.0)
+    masked_pixel_percentage: float = Field(default=0.2, ge=0.05, le=10.0)
     strategy: Literal["uniform", "median"] = Field(default="uniform")
     struct_mask_axis: Literal["horizontal", "vertical", "none"] = Field(default="none")
     struct_mask_span: int = Field(default=5, ge=3, le=15)

careamics/config/vae_algorithm_model.py

@@ -0,0 +1,171 @@
+"""Algorithm configuration."""
+
+from __future__ import annotations
+
+from pprint import pformat
+from typing import Literal, Optional, Union
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+from typing_extensions import Self
+
+from careamics.config.support import SupportedAlgorithm, SupportedLoss
+
+from .architectures import CustomModel, LVAEModel
+from .likelihood_model import GaussianLikelihoodConfig, NMLikelihoodConfig
+from .nm_model import MultiChannelNMConfig
+from .optimizer_models import LrSchedulerModel, OptimizerModel
+
+
+class VAEAlgorithmConfig(BaseModel):
+    """Algorithm configuration.
+
+    This Pydantic model validates the parameters governing the components of the
+    training algorithm: which algorithm, loss function, model architecture, optimizer,
+    and learning rate scheduler to use.
+
+    Currently, we only support N2V, CARE, N2N and custom models. The `n2v` algorithm is
+    only compatible with `n2v` loss and `UNet` architecture. The `custom` algorithm
+    allows you to register your own architecture and select it using its name as
+    `name` in the custom pydantic model.
+
+    Attributes
+    ----------
+    algorithm : algorithm: Literal["musplit", "denoisplit", "custom"]
+        Algorithm to use.
+    loss : Literal["musplit", "denoisplit", "denoisplit_musplit"]
+        Loss function to use.
+    model : Union[LVAEModel, CustomModel]
+        Model architecture to use.
+    noise_model: Optional[MultiChannelNmModel]
+        Noise model to use.
+    noise_model_likelihood_model: Optional[NMLikelihoodModel]
+        Noise model likelihood model to use.
+    gaussian_likelihood_model: Optional[GaussianLikelihoodModel]
+        Gaussian likelihood model to use.
+    optimizer : OptimizerModel, optional
+        Optimizer to use.
+    lr_scheduler : LrSchedulerModel, optional
+        Learning rate scheduler to use.
+
+    Raises
+    ------
+    ValueError
+        Algorithm parameter type validation errors.
+    ValueError
+        If the algorithm, loss and model are not compatible.
+
+    Examples
+    --------
+    # TODO add once finalized
+    """
+
+    # Pydantic class configuration
+    model_config = ConfigDict(
+        protected_namespaces=(),  # allows to use model_* as a field name
+        validate_assignment=True,
+        extra="allow",
+    )
+
+    # Mandatory fields
+    # defined in SupportedAlgorithm
+    # TODO: Use supported Enum classes for typing?
+    #   - values can still be passed as strings and they will be cast to Enum
+    algorithm_type: Literal["vae"]
+    algorithm: Literal["musplit", "denoisplit", "custom"]
+    loss: Literal["musplit", "denoisplit", "denoisplit_musplit"]
+    model: Union[LVAEModel, CustomModel] = Field(discriminator="architecture")
+
+    # TODO: these are configs, change naming of attrs
+    noise_model: Optional[MultiChannelNMConfig] = None
+    noise_model_likelihood_model: Optional[NMLikelihoodConfig] = None
+    gaussian_likelihood_model: Optional[GaussianLikelihoodConfig] = None
+
+    # Optional fields
+    optimizer: OptimizerModel = OptimizerModel()
+    """Optimizer to use, defined in SupportedOptimizer."""
+
+    lr_scheduler: LrSchedulerModel = LrSchedulerModel()
+
+    @model_validator(mode="after")
+    def algorithm_cross_validation(self: Self) -> Self:
+        """Validate the algorithm model based on `algorithm`.
+
+        Returns
+        -------
+        Self
+            The validated model.
+        """
+        # musplit
+        if self.algorithm == SupportedAlgorithm.MUSPLIT:
+            if self.loss != SupportedLoss.MUSPLIT:
+                raise ValueError(
+                    f"Algorithm {self.algorithm} only supports loss `musplit`."
+                )
+
+        if self.algorithm == SupportedAlgorithm.DENOISPLIT:
+            if self.loss not in [
+                SupportedLoss.DENOISPLIT,
+                SupportedLoss.DENOISPLIT_MUSPLIT,
+            ]:
+                raise ValueError(
+                    f"Algorithm {self.algorithm} only supports loss `denoisplit` "
+                    "or `denoisplit_musplit."
+                )
+            if (
+                self.loss == SupportedLoss.DENOISPLIT
+                and self.model.predict_logvar is not None
+            ):
+                raise ValueError(
+                    "Algorithm `denoisplit` with loss `denoisplit` only supports "
+                    "`predict_logvar` as `None`."
+                )
+            if self.noise_model is None:
+                raise ValueError("Algorithm `denoisplit` requires a noise model.")
+        # TODO: what if algorithm is not musplit or denoisplit (HDN?)
+        return self
+
+    @model_validator(mode="after")
+    def output_channels_validation(self: Self) -> Self:
+        """Validate the consistency between number of out channels and noise models.
+
+        Returns
+        -------
+        Self
+            The validated model.
+        """
+        if self.noise_model is not None:
+            assert self.model.output_channels == len(self.noise_model.noise_models), (
+                f"Number of output channels ({self.model.output_channels}) must match "
+                f"the number of noise models ({len(self.noise_model.noise_models)})."
+            )
+        return self
+
+    @model_validator(mode="after")
+    def predict_logvar_validation(self: Self) -> Self:
+        """Validate the consistency of `predict_logvar` throughout the model.
+
+        Returns
+        -------
+        Self
+            The validated model.
+        """
+        if self.gaussian_likelihood_model is not None:
+            assert (
+                self.model.predict_logvar
+                == self.gaussian_likelihood_model.predict_logvar
+            ), (
+                f"Model `predict_logvar` ({self.model.predict_logvar}) must match "
+                "Gaussian likelihood model `predict_logvar` "
+                f"({self.gaussian_likelihood_model.predict_logvar}).",
+            )
+        return self
+
+    def __str__(self) -> str:
+        """Pretty string representing the configuration.
+
+        Returns
+        -------
+        str
+            Pretty string.
+        """
+        return pformat(self.model_dump())

careamics/dataset/tiling/lvae_tiled_patching.py

@@ -0,0 +1,282 @@
+"""Functions to reimplement the tiling in the Disentangle repository."""
+
+import builtins
+import itertools
+from typing import Any, Generator, Optional, Union
+
+import numpy as np
+from numpy.typing import NDArray
+
+from careamics.config.tile_information import TileInformation
+
+
+def extract_tiles(
+    arr: NDArray,
+    tile_size: NDArray[np.int_],
+    overlaps: NDArray[np.int_],
+    padding_kwargs: Optional[dict[str, Any]] = None,
+) -> Generator[tuple[NDArray, TileInformation], None, None]:
+    """Generate tiles from the input array with specified overlap.
+
+    The tiles cover the whole array; which will be additionally padded, to ensure that
+    the section of the tile that contributes to the final image comes from the center
+    of the tile.
+
+    The method returns a generator that yields tuples of array and tile information,
+    the latter includes whether the tile is the last one, the coordinates of the
+    overlap crop, and the coordinates of the stitched tile.
+
+    Input array should have shape SC(Z)YX, while the returned tiles have shape C(Z)YX,
+    where C can be a singleton.
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Array of shape (S, C, (Z), Y, X).
+    tile_size : 1D numpy.ndarray of tuple
+        Tile sizes in each dimension, of length 2 or 3.
+    overlaps : 1D numpy.ndarray of tuple
+        Overlap values in each dimension, of length 2 or 3.
+    padding_kwargs : dict, optional
+        The arguments of `np.pad` after the first two arguments, `array` and
+        `pad_width`. If not specified the default will be `{"mode": "reflect"}`. See
+        `numpy.pad` docs:
+        https://numpy.org/doc/stable/reference/generated/numpy.pad.html.
+
+    Yields
+    ------
+    Generator[Tuple[np.ndarray, TileInformation], None, None]
+        Tile generator, yields the tile and additional information.
+    """
+    if padding_kwargs is None:
+        padding_kwargs = {"mode": "reflect"}
+
+    # Iterate over num samples (S)
+    for sample_idx in range(arr.shape[0]):
+        sample = arr[sample_idx, ...]
+        data_shape = np.array(sample.shape)
+
+        # add padding to ensure evenly spaced & overlapping tiles.
+        spatial_padding = compute_padding(data_shape, tile_size, overlaps)
+        padding = ((0, 0), *spatial_padding)
+        sample = np.pad(sample, padding, **padding_kwargs)
+
+        # The number of tiles in each dimension, should be of length 2 or 3
+        tile_grid_shape = compute_tile_grid_shape(data_shape, tile_size, overlaps)
+        # itertools.product is equivalent of nested loops
+
+        stitch_size = tile_size - overlaps
+        for tile_grid_coords in itertools.product(*[range(n) for n in tile_grid_shape]):
+
+            # calculate crop coordinates
+            crop_coords_start = np.array(tile_grid_coords) * stitch_size
+            crop_slices: tuple[Union[builtins.ellipsis, slice], ...] = (
+                ...,
+                *[
+                    slice(coords, coords + extent)
+                    for coords, extent in zip(crop_coords_start, tile_size)
+                ],
+            )
+            tile = sample[crop_slices]
+
+            tile_info = compute_tile_info(
+                np.array(tile_grid_coords),
+                np.array(data_shape),
+                np.array(tile_size),
+                np.array(overlaps),
+                sample_idx,
+            )
+            # TODO: kinda weird this is a generator,
+            #   -> doesn't really save memory ? Don't think there are any places the
+            #    tiles are not exracted all at the same time.
+            #   Although I guess it would make sense for a zarr tile extractor.
+            yield tile, tile_info
+
+
+def compute_tile_info(
+    tile_grid_coords: NDArray[np.int_],
+    data_shape: NDArray[np.int_],
+    tile_size: NDArray[np.int_],
+    overlaps: NDArray[np.int_],
+    sample_id: int = 0,
+) -> TileInformation:
+    """
+    Compute the tile information for a tile with the coordinates `tile_grid_coords`.
+
+    Parameters
+    ----------
+    tile_grid_coords : 1D np.array of int
+        The coordinates of the tile within the tile grid, ((Z), Y, X), i.e. for 2D
+        tiling the coordinates for the second tile in the first row of tiles would be
+        (0, 1).
+    data_shape : 1D np.array of int
+        The shape of the data, should be (C, (Z), Y, X) where Z is optional.
+    tile_size : 1D np.array of int
+        Tile sizes in each dimension, of length 2 or 3.
+    overlaps : 1D np.array of int
+        Overlap values in each dimension, of length 2 or 3.
+    sample_id : int, default=0
+        An ID to identify which sample a tile belongs to.
+
+    Returns
+    -------
+    TileInformation
+        Information that describes how to crop and stitch a tile to create a full image.
+    """
+    spatial_dims_shape = data_shape[-len(tile_size) :]
+
+    # The extent of the tile which will make up part of the stitched image.
+    stitch_size = tile_size - overlaps
+    stitch_coords_start = tile_grid_coords * stitch_size
+    stitch_coords_end = stitch_coords_start + stitch_size
+
+    tile_coords_start = stitch_coords_start - overlaps // 2
+
+    # --- replace out of bounds indices
+    out_of_lower_bound = stitch_coords_start < 0
+    out_of_upper_bound = stitch_coords_end > spatial_dims_shape
+    stitch_coords_start[out_of_lower_bound] = 0
+    stitch_coords_end[out_of_upper_bound] = spatial_dims_shape[out_of_upper_bound]
+
+    # --- calculate overlap crop coords
+    overlap_crop_coords_start = stitch_coords_start - tile_coords_start
+    overlap_crop_coords_end = overlap_crop_coords_start + (
+        stitch_coords_end - stitch_coords_start
+    )
+
+    # --- combine start and end
+    stitch_coords = tuple(
+        (start, end) for start, end in zip(stitch_coords_start, stitch_coords_end)
+    )
+    overlap_crop_coords = tuple(
+        (start, end)
+        for start, end in zip(overlap_crop_coords_start, overlap_crop_coords_end)
+    )
+
+    # --- Check if last tile
+    tile_grid_shape = np.array(compute_tile_grid_shape(data_shape, tile_size, overlaps))
+    last_tile = (tile_grid_coords == (tile_grid_shape - 1)).all()
+
+    tile_info = TileInformation(
+        array_shape=data_shape,
+        last_tile=last_tile,
+        overlap_crop_coords=overlap_crop_coords,
+        stitch_coords=stitch_coords,
+        sample_id=sample_id,
+    )
+    return tile_info
+
+
+def compute_padding(
+    data_shape: NDArray[np.int_],
+    tile_size: NDArray[np.int_],
+    overlaps: NDArray[np.int_],
+) -> tuple[tuple[int, int], ...]:
+    """
+    Calculate padding to ensure stitched data comes from the center of a tile.
+
+    Padding is added to an array with shape `data_shape` so that when tiles are
+    stitched together, the data used always comes from the center of a tile, even for
+    tiles at the boundaries of the array.
+
+    Parameters
+    ----------
+    data_shape : 1D numpy.array of int
+        The shape of the data to be tiled and stitched together, (S, C, (Z), Y, X).
+    tile_size : 1D numpy.array of int
+        The tile size in each dimension, ((Z), Y, X).
+    overlaps : 1D numpy.array of int
+        The tile overlap in each dimension, ((Z), Y, X).
+
+    Returns
+    -------
+    tuple of (int, int)
+        A tuple specifying the padding to add in each dimension, each element is a two
+        element tuple specifying the padding to add before and after the data. This
+        can be used as the `pad_width` argument to `numpy.pad`.
+    """
+    tile_grid_shape = np.array(compute_tile_grid_shape(data_shape, tile_size, overlaps))
+    covered_shape = (tile_size - overlaps) * tile_grid_shape + overlaps
+
+    pad_before = overlaps // 2
+    pad_after = covered_shape - data_shape[-len(tile_size) :] - pad_before
+
+    return tuple((before, after) for before, after in zip(pad_before, pad_after))
+
+
+def n_tiles_1d(axis_size: int, tile_size: int, overlap: int) -> int:
+    """Calculate the number of tiles in a specific dimension.
+
+    Parameters
+    ----------
+    axis_size : int
+        The length of the data for in a specific dimension.
+    tile_size : int
+        The length of the tiles in a specific dimension.
+    overlap : int
+        The tile overlap in a specific dimension.
+
+    Returns
+    -------
+    int
+        The number of tiles that fit in one dimension given the arguments.
+    """
+    return int(np.ceil(axis_size / (tile_size - overlap)))
+
+
+def total_n_tiles(
+    data_shape: tuple[int, ...], tile_size: tuple[int, ...], overlaps: tuple[int, ...]
+) -> int:
+    """Calculate The total number of tiles over all dimensions.
+
+    Parameters
+    ----------
+    data_shape : 1D numpy.array of int
+        The shape of the data to be tiled and stitched together, (S, C, (Z), Y, X).
+    tile_size : 1D numpy.array of int
+        The tile size in each dimension, ((Z), Y, X).
+    overlaps : 1D numpy.array of int
+        The tile overlap in each dimension, ((Z), Y, X).
+
+
+    Returns
+    -------
+    int
+        The total number of tiles over all dimensions.
+    """
+    result = 1
+    # assume spatial dimension are the last dimensions so iterate backwards
+    for i in range(-1, -len(tile_size) - 1, -1):
+        result = result * n_tiles_1d(data_shape[i], tile_size[i], overlaps[i])
+
+    return result
+
+
+def compute_tile_grid_shape(
+    data_shape: NDArray[np.int_],
+    tile_size: NDArray[np.int_],
+    overlaps: NDArray[np.int_],
+) -> tuple[int, ...]:
+    """Calculate the number of tiles in each dimension.
+
+    This can be thought of as a grid of tiles.
+
+    Parameters
+    ----------
+    data_shape : 1D numpy.array of int
+        The shape of the data to be tiled and stitched together, (S, C, (Z), Y, X).
+    tile_size : 1D numpy.array of int
+        The tile size in each dimension, ((Z), Y, X).
+    overlaps : 1D numpy.array of int
+        The tile overlap in each dimension, ((Z), Y, X).
+
+    Returns
+    -------
+    tuple of int
+        The number of tiles in each direction, ((Z, Y, X)).
+    """
+    shape = [0 for _ in range(len(tile_size))]
+    # assume spatial dimension are the last dimensions so iterate backwards
+    for i in range(-1, -len(tile_size) - 1, -1):
+        shape[i] = n_tiles_1d(data_shape[i], tile_size[i], overlaps[i])
+    return tuple(shape)