careamics 0.0.12__py3-none-any.whl → 0.0.13__py3-none-any.whl

careamics/config/configuration_factories.py

@@ -1,12 +1,13 @@
 """Convenience functions to create configurations for training and inference."""
 
+from collections.abc import Sequence
 from typing import Annotated, Any, Literal, Optional, Union
 
 from pydantic import Field, TypeAdapter
 
 from careamics.config.algorithms import CAREAlgorithm, N2NAlgorithm, N2VAlgorithm
 from careamics.config.architectures import UNetModel
-from careamics.config.data import DataConfig
+from careamics.config.data import DataConfig, NGDataConfig
 from careamics.config.support import (
     SupportedArchitecture,
     SupportedPixelManipulation,
@@ -24,7 +25,7 @@ from .configuration import Configuration
 
 
 def algorithm_factory(
-    algorithm: dict[str, Any]
+    algorithm: dict[str, Any],
 ) -> Union[N2VAlgorithm, N2NAlgorithm, CAREAlgorithm]:
     """
     Create an algorithm model for training CAREamics.
@@ -49,7 +50,7 @@ def algorithm_factory(
 
 
 def _list_spatial_augmentations(
-    augmentations: Optional[list[SPATIAL_TRANSFORMS_UNION]],
+    augmentations: Optional[list[SPATIAL_TRANSFORMS_UNION]] = None,
 ) -> list[SPATIAL_TRANSFORMS_UNION]:
     """
     List the augmentations to apply.
@@ -153,6 +154,10 @@ def _create_algorithm_configuration(
     n_channels_out: int,
     use_n2v2: bool = False,
     model_params: Optional[dict] = None,
+    optimizer: Literal["Adam", "Adamax", "SGD"] = "Adam",
+    optimizer_params: Optional[dict[str, Any]] = None,
+    lr_scheduler: Literal["ReduceLROnPlateau", "StepLR"] = "ReduceLROnPlateau",
+    lr_scheduler_params: Optional[dict[str, Any]] = None,
 ) -> dict:
     """
     Create a dictionary with the parameters of the algorithm model.
@@ -171,10 +176,20 @@ def _create_algorithm_configuration(
         Number of input channels.
     n_channels_out : int
         Number of output channels.
-    use_n2v2 : bool, optional
-        Whether to use N2V2, by default False.
-    model_params : dict
+    use_n2v2 : bool, default=false
+        Whether to use N2V2.
+    model_params : dict, default=None
         UNetModel parameters.
+    optimizer : {"Adam", "Adamax", "SGD"}, default="Adam"
+        Optimizer to use.
+    optimizer_params : dict, default=None
+        Parameters for the optimizer, see PyTorch documentation for more details.
+    lr_scheduler : {"ReduceLROnPlateau", "StepLR"}, default="ReduceLROnPlateau"
+        Learning rate scheduler to use.
+    lr_scheduler_params : dict, default=None
+        Parameters for the learning rate scheduler, see PyTorch documentation for more
+        details.
+
 
     Returns
     -------
@@ -195,11 +210,19 @@ def _create_algorithm_configuration(
         "algorithm": algorithm,
         "loss": loss,
         "model": unet_model,
+        "optimizer": {
+            "name": optimizer,
+            "parameters": {} if optimizer_params is None else optimizer_params,
+        },
+        "lr_scheduler": {
+            "name": lr_scheduler,
+            "parameters": {} if lr_scheduler_params is None else lr_scheduler_params,
+        },
     }
 
 
 def _create_data_configuration(
-    data_type: Literal["array", "tiff", "custom"],
+    data_type: Literal["array", "tiff", "czi", "custom"],
     axes: str,
     patch_size: list[int],
     batch_size: int,
@@ -212,7 +235,7 @@ def _create_data_configuration(
 
     Parameters
     ----------
-    data_type : {"array", "tiff", "custom"}
+    data_type : {"array", "tiff", "czi", "custom"}
         Type of the data.
     axes : str
         Axes of the data.
@@ -254,8 +277,89 @@ def _create_data_configuration(
     return DataConfig(**data)
 
 
+def _create_ng_data_configuration(
+    data_type: Literal["array", "tiff", "custom"],
+    axes: str,
+    patch_size: Sequence[int],
+    batch_size: int,
+    augmentations: list[SPATIAL_TRANSFORMS_UNION],
+    patch_overlaps: Optional[Sequence[int]] = None,
+    train_dataloader_params: Optional[dict[str, Any]] = None,
+    val_dataloader_params: Optional[dict[str, Any]] = None,
+    test_dataloader_params: Optional[dict[str, Any]] = None,
+    seed: Optional[int] = None,
+) -> NGDataConfig:
+    """
+    Create a dictionary with the parameters of the data model.
+
+    Parameters
+    ----------
+    data_type : {"array", "tiff", "custom"}
+        Type of the data.
+    axes : str
+        Axes of the data.
+    patch_size : list of int
+        Size of the patches along the spatial dimensions.
+    batch_size : int
+        Batch size.
+    augmentations : list of transforms
+        List of transforms to apply.
+    patch_overlaps : Sequence of int, default=None
+        Overlaps between patches in each spatial dimension, only used with "sequential"
+        patching. If `None`, no overlap is applied. The overlap must be smaller than
+        the patch size in each spatial dimension, and the number of dimensions be either
+        2 or 3.
+    train_dataloader_params : dict
+        Parameters for the training dataloader, see PyTorch notes, by default None.
+    val_dataloader_params : dict
+        Parameters for the validation dataloader, see PyTorch notes, by default None.
+    test_dataloader_params : dict
+        Parameters for the test dataloader, see PyTorch notes, by default None.
+    seed : int, default=None
+        Random seed for reproducibility. If `None`, no seed is set.
+
+    Returns
+    -------
+    NGDataConfig
+        Next-Generation Data model with the specified parameters.
+    """
+    # data model
+    data = {
+        "data_type": data_type,
+        "axes": axes,
+        "batch_size": batch_size,
+        "transforms": augmentations,
+        "seed": seed,
+    }
+    # don't override defaults set in DataConfig class
+    if train_dataloader_params is not None:
+        # the presence of `shuffle` key in the dataloader parameters is enforced
+        # by the NGDataConfig class
+        if "shuffle" not in train_dataloader_params:
+            train_dataloader_params["shuffle"] = True
+
+        data["train_dataloader_params"] = train_dataloader_params
+
+    if val_dataloader_params is not None:
+        data["val_dataloader_params"] = val_dataloader_params
+
+    if test_dataloader_params is not None:
+        data["test_dataloader_params"] = test_dataloader_params
+
+    # add training patching
+    data["patching"] = {
+        "name": "random",
+        "patch_size": patch_size,
+        "overlaps": patch_overlaps,
+    }
+
+    return NGDataConfig(**data)
+
+
 def _create_training_configuration(
-    num_epochs: int, logger: Literal["wandb", "tensorboard", "none"]
+    num_epochs: int,
+    logger: Literal["wandb", "tensorboard", "none"],
+    checkpoint_params: Optional[dict[str, Any]] = None,
 ) -> TrainingConfig:
     """
     Create a dictionary with the parameters of the training model.
@@ -266,6 +370,9 @@ def _create_training_configuration(
         Number of epochs.
     logger : {"wandb", "tensorboard", "none"}
         Logger to use.
+    checkpoint_params : dict, default=None
+        Parameters for the checkpoint callback, see PyTorch Lightning documentation
+        (`ModelCheckpoint`) for the list of available parameters.
 
     Returns
     -------
@@ -275,6 +382,7 @@ def _create_training_configuration(
     return TrainingConfig(
         num_epochs=num_epochs,
         logger=None if logger == "none" else logger,
+        checkpoint_callback={} if checkpoint_params is None else checkpoint_params,
     )
 
 
@@ -282,7 +390,7 @@ def _create_training_configuration(
 def _create_supervised_config_dict(
     algorithm: Literal["care", "n2n"],
     experiment_name: str,
-    data_type: Literal["array", "tiff", "custom"],
+    data_type: Literal["array", "tiff", "czi", "custom"],
     axes: str,
     patch_size: list[int],
     batch_size: int,
@@ -294,8 +402,13 @@ def _create_supervised_config_dict(
     n_channels_out: Optional[int] = None,
     logger: Literal["wandb", "tensorboard", "none"] = "none",
     model_params: Optional[dict] = None,
+    optimizer: Literal["Adam", "Adamax", "SGD"] = "Adam",
+    optimizer_params: Optional[dict[str, Any]] = None,
+    lr_scheduler: Literal["ReduceLROnPlateau", "StepLR"] = "ReduceLROnPlateau",
+    lr_scheduler_params: Optional[dict[str, Any]] = None,
     train_dataloader_params: Optional[dict[str, Any]] = None,
     val_dataloader_params: Optional[dict[str, Any]] = None,
+    checkpoint_params: Optional[dict[str, Any]] = None,
 ) -> dict:
     """
     Create a configuration for training CARE or Noise2Noise.
@@ -306,7 +419,7 @@ def _create_supervised_config_dict(
         Algorithm to use.
     experiment_name : str
         Name of the experiment.
-    data_type : Literal["array", "tiff", "custom"]
+    data_type : Literal["array", "tiff", "czi", "custom"]
         Type of the data.
     axes : str
         Axes of the data (e.g. SYX).
@@ -330,12 +443,24 @@ def _create_supervised_config_dict(
         Number of channels out.
     logger : Literal["wandb", "tensorboard", "none"], optional
         Logger to use, by default "none".
-    model_params : dict, optional
-        UNetModel parameters, by default {}.
+    model_params : dict, default=None
+        UNetModel parameters.
+    optimizer : {"Adam", "Adamax", "SGD"}, default="Adam"
+        Optimizer to use.
+    optimizer_params : dict, default=None
+        Parameters for the optimizer, see PyTorch documentation for more details.
+    lr_scheduler : {"ReduceLROnPlateau", "StepLR"}, default="ReduceLROnPlateau"
+        Learning rate scheduler to use.
+    lr_scheduler_params : dict, default=None
+        Parameters for the learning rate scheduler, see PyTorch documentation for more
+        details.
     train_dataloader_params : dict
         Parameters for the training dataloader, see PyTorch notes, by default None.
     val_dataloader_params : dict
         Parameters for the validation dataloader, see PyTorch notes, by default None.
+    checkpoint_params : dict, default=None
+        Parameters for the checkpoint callback, see PyTorch Lightning documentation
+        (`ModelCheckpoint`) for the list of available parameters.
 
     Returns
     -------
@@ -376,6 +501,10 @@ def _create_supervised_config_dict(
         n_channels_in=n_channels_in,
         n_channels_out=n_channels_out,
         model_params=model_params,
+        optimizer=optimizer,
+        optimizer_params=optimizer_params,
+        lr_scheduler=lr_scheduler,
+        lr_scheduler_params=lr_scheduler_params,
     )
 
     # data
@@ -393,6 +522,7 @@ def _create_supervised_config_dict(
     training_params = _create_training_configuration(
         num_epochs=num_epochs,
         logger=logger,
+        checkpoint_params=checkpoint_params,
     )
 
     return {
@@ -405,7 +535,7 @@ def _create_supervised_config_dict(
 
 def create_care_configuration(
     experiment_name: str,
-    data_type: Literal["array", "tiff", "custom"],
+    data_type: Literal["array", "tiff", "czi", "custom"],
     axes: str,
     patch_size: list[int],
     batch_size: int,
@@ -417,8 +547,13 @@ def create_care_configuration(
     n_channels_out: Optional[int] = None,
     logger: Literal["wandb", "tensorboard", "none"] = "none",
     model_params: Optional[dict] = None,
+    optimizer: Literal["Adam", "Adamax", "SGD"] = "Adam",
+    optimizer_params: Optional[dict[str, Any]] = None,
+    lr_scheduler: Literal["ReduceLROnPlateau", "StepLR"] = "ReduceLROnPlateau",
+    lr_scheduler_params: Optional[dict[str, Any]] = None,
     train_dataloader_params: Optional[dict[str, Any]] = None,
     val_dataloader_params: Optional[dict[str, Any]] = None,
+    checkpoint_params: Optional[dict[str, Any]] = None,
 ) -> Configuration:
     """
     Create a configuration for training CARE.
@@ -445,7 +580,7 @@ def create_care_configuration(
     ----------
     experiment_name : str
         Name of the experiment.
-    data_type : Literal["array", "tiff", "custom"]
+    data_type : Literal["array", "tiff", "czi", "custom"]
         Type of the data.
     axes : str
         Axes of the data (e.g. SYX).
@@ -471,6 +606,15 @@ def create_care_configuration(
         Logger to use.
     model_params : dict, default=None
         UNetModel parameters.
+    optimizer : Literal["Adam", "Adamax", "SGD"], default="Adam"
+        Optimizer to use.
+    optimizer_params : dict, default=None
+        Parameters for the optimizer, see PyTorch documentation for more details.
+    lr_scheduler : Literal["ReduceLROnPlateau", "StepLR"], default="ReduceLROnPlateau"
+        Learning rate scheduler to use.
+    lr_scheduler_params : dict, default=None
+        Parameters for the learning rate scheduler, see PyTorch documentation for more
+        details.
     train_dataloader_params : dict, optional
         Parameters for the training dataloader, see the PyTorch docs for `DataLoader`.
         If left as `None`, the dict `{"shuffle": True}` will be used, this is set in
@@ -479,6 +623,9 @@ def create_care_configuration(
         Parameters for the validation dataloader, see PyTorch the docs for `DataLoader`.
         If left as `None`, the empty dict `{}` will be used, this is set in the
         `GeneralDataConfig`.
+    checkpoint_params : dict, default=None
+        Parameters for the checkpoint callback, see PyTorch Lightning documentation
+        (`ModelCheckpoint`) for the list of available parameters.
 
     Returns
     -------
@@ -551,6 +698,29 @@ def create_care_configuration(
     ...     n_channels_in=3,
     ...     n_channels_out=1 # if applicable
     ... )
+
+    If you would like to train on CZI files, use `"czi"` as `data_type` and `"SCYX"` as
+    `axes` for 2-D or `"SCZYX"` for 3-D denoising. Note that `"SCYX"` can also be used
+    for 3-D data but spatial context along the Z dimension will then not be taken into
+    account.
+    >>> config_2d = create_care_configuration(
+    ...     experiment_name="care_experiment",
+    ...     data_type="czi",
+    ...     axes="SCYX",
+    ...     patch_size=[64, 64],
+    ...     batch_size=32,
+    ...     num_epochs=100,
+    ...     n_channels_in=1,
+    ... )
+    >>> config_3d = create_care_configuration(
+    ...     experiment_name="care_experiment",
+    ...     data_type="czi",
+    ...     axes="SCZYX",
+    ...     patch_size=[16, 64, 64],
+    ...     batch_size=16,
+    ...     num_epochs=100,
+    ...     n_channels_in=1,
+    ... )
     """
     return Configuration(
         **_create_supervised_config_dict(
@@ -568,15 +738,20 @@ def create_care_configuration(
             n_channels_out=n_channels_out,
             logger=logger,
             model_params=model_params,
+            optimizer=optimizer,
+            optimizer_params=optimizer_params,
+            lr_scheduler=lr_scheduler,
+            lr_scheduler_params=lr_scheduler_params,
             train_dataloader_params=train_dataloader_params,
             val_dataloader_params=val_dataloader_params,
+            checkpoint_params=checkpoint_params,
         )
     )
 
 
 def create_n2n_configuration(
     experiment_name: str,
-    data_type: Literal["array", "tiff", "custom"],
+    data_type: Literal["array", "tiff", "czi", "custom"],
     axes: str,
     patch_size: list[int],
     batch_size: int,
@@ -588,8 +763,13 @@ def create_n2n_configuration(
     n_channels_out: Optional[int] = None,
     logger: Literal["wandb", "tensorboard", "none"] = "none",
     model_params: Optional[dict] = None,
+    optimizer: Literal["Adam", "Adamax", "SGD"] = "Adam",
+    optimizer_params: Optional[dict[str, Any]] = None,
+    lr_scheduler: Literal["ReduceLROnPlateau", "StepLR"] = "ReduceLROnPlateau",
+    lr_scheduler_params: Optional[dict[str, Any]] = None,
     train_dataloader_params: Optional[dict[str, Any]] = None,
     val_dataloader_params: Optional[dict[str, Any]] = None,
+    checkpoint_params: Optional[dict[str, Any]] = None,
 ) -> Configuration:
     """
     Create a configuration for training Noise2Noise.
@@ -616,7 +796,7 @@ def create_n2n_configuration(
     ----------
     experiment_name : str
         Name of the experiment.
-    data_type : Literal["array", "tiff", "custom"]
+    data_type : Literal["array", "tiff", "czi", "custom"]
         Type of the data.
     axes : str
         Axes of the data (e.g. SYX).
@@ -640,8 +820,17 @@ def create_n2n_configuration(
         Number of channels out.
     logger : Literal["wandb", "tensorboard", "none"], optional
         Logger to use, by default "none".
-    model_params : dict, optional
-        UNetModel parameters, by default {}.
+    model_params : dict, default=None
+        UNetModel parameters.
+    optimizer : Literal["Adam", "Adamax", "SGD"], default="Adam"
+        Optimizer to use.
+    optimizer_params : dict, default=None
+        Parameters for the optimizer, see PyTorch documentation for more details.
+    lr_scheduler : Literal["ReduceLROnPlateau", "StepLR"], default="ReduceLROnPlateau"
+        Learning rate scheduler to use.
+    lr_scheduler_params : dict, default=None
+        Parameters for the learning rate scheduler, see PyTorch documentation for more
+        details.
     train_dataloader_params : dict, optional
         Parameters for the training dataloader, see the PyTorch docs for `DataLoader`.
         If left as `None`, the dict `{"shuffle": True}` will be used, this is set in
@@ -650,6 +839,9 @@ def create_n2n_configuration(
         Parameters for the validation dataloader, see PyTorch the docs for `DataLoader`.
         If left as `None`, the empty dict `{}` will be used, this is set in the
         `GeneralDataConfig`.
+    checkpoint_params : dict, default=None
+        Parameters for the checkpoint callback, see PyTorch Lightning documentation
+        (`ModelCheckpoint`) for the list of available parameters.
 
     Returns
     -------
@@ -722,6 +914,29 @@ def create_n2n_configuration(
     ...     n_channels_in=3,
     ...     n_channels_out=1 # if applicable
     ... )
+
+    If you would like to train on CZI files, use `"czi"` as `data_type` and `"SCYX"` as
+    `axes` for 2-D or `"SCZYX"` for 3-D denoising. Note that `"SCYX"` can also be used
+    for 3-D data but spatial context along the Z dimension will then not be taken into
+    account.
+    >>> config_2d = create_n2n_configuration(
+    ...     experiment_name="n2n_experiment",
+    ...     data_type="czi",
+    ...     axes="SCYX",
+    ...     patch_size=[64, 64],
+    ...     batch_size=32,
+    ...     num_epochs=100,
+    ...     n_channels_in=1,
+    ... )
+    >>> config_3d = create_n2n_configuration(
+    ...     experiment_name="n2n_experiment",
+    ...     data_type="czi",
+    ...     axes="SCZYX",
+    ...     patch_size=[16, 64, 64],
+    ...     batch_size=16,
+    ...     num_epochs=100,
+    ...     n_channels_in=1,
+    ... )
     """
     return Configuration(
         **_create_supervised_config_dict(
@@ -739,15 +954,20 @@ def create_n2n_configuration(
             n_channels_out=n_channels_out,
             logger=logger,
             model_params=model_params,
+            optimizer=optimizer,
+            optimizer_params=optimizer_params,
+            lr_scheduler=lr_scheduler,
+            lr_scheduler_params=lr_scheduler_params,
             train_dataloader_params=train_dataloader_params,
             val_dataloader_params=val_dataloader_params,
+            checkpoint_params=checkpoint_params,
         )
     )
 
 
 def create_n2v_configuration(
     experiment_name: str,
-    data_type: Literal["array", "tiff", "custom"],
+    data_type: Literal["array", "tiff", "czi", "custom"],
     axes: str,
     patch_size: list[int],
     batch_size: int,
@@ -762,8 +982,13 @@ def create_n2v_configuration(
     struct_n2v_span: int = 5,
     logger: Literal["wandb", "tensorboard", "none"] = "none",
     model_params: Optional[dict] = None,
+    optimizer: Literal["Adam", "Adamax", "SGD"] = "Adam",
+    optimizer_params: Optional[dict[str, Any]] = None,
+    lr_scheduler: Literal["ReduceLROnPlateau", "StepLR"] = "ReduceLROnPlateau",
+    lr_scheduler_params: Optional[dict[str, Any]] = None,
     train_dataloader_params: Optional[dict[str, Any]] = None,
     val_dataloader_params: Optional[dict[str, Any]] = None,
+    checkpoint_params: Optional[dict[str, Any]] = None,
 ) -> Configuration:
     """
     Create a configuration for training Noise2Void.
@@ -810,7 +1035,7 @@ def create_n2v_configuration(
     ----------
     experiment_name : str
         Name of the experiment.
-    data_type : Literal["array", "tiff", "custom"]
+    data_type : Literal["array", "tiff", "czi", "custom"]
         Type of the data.
     axes : str
         Axes of the data (e.g. SYX).
@@ -840,8 +1065,17 @@ def create_n2v_configuration(
         Span of the structN2V mask, by default 5.
     logger : Literal["wandb", "tensorboard", "none"], optional
         Logger to use, by default "none".
-    model_params : dict, optional
-        UNetModel parameters, by default None.
+    model_params : dict, default=None
+        UNetModel parameters.
+    optimizer : Literal["Adam", "Adamax", "SGD"], default="Adam"
+        Optimizer to use.
+    optimizer_params : dict, default=None
+        Parameters for the optimizer, see PyTorch documentation for more details.
+    lr_scheduler : Literal["ReduceLROnPlateau", "StepLR"], default="ReduceLROnPlateau"
+        Learning rate scheduler to use.
+    lr_scheduler_params : dict, default=None
+        Parameters for the learning rate scheduler, see PyTorch documentation for more
+        details.
     train_dataloader_params : dict, optional
         Parameters for the training dataloader, see the PyTorch docs for `DataLoader`.
         If left as `None`, the dict `{"shuffle": True}` will be used, this is set in
@@ -850,6 +1084,9 @@ def create_n2v_configuration(
         Parameters for the validation dataloader, see PyTorch the docs for `DataLoader`.
         If left as `None`, the empty dict `{}` will be used, this is set in the
         `GeneralDataConfig`.
+    checkpoint_params : dict, default=None
+        Parameters for the checkpoint callback, see PyTorch Lightning documentation
+        (`ModelCheckpoint`) for the list of available parameters.
 
     Returns
     -------
@@ -942,6 +1179,29 @@ def create_n2v_configuration(
     ...     independent_channels=False,
     ...     n_channels=3
     ... )
+
+    If you would like to train on CZI files, use `"czi"` as `data_type` and `"SCYX"` as
+    `axes` for 2-D or `"SCZYX"` for 3-D denoising. Note that `"SCYX"` can also be used
+    for 3-D data but spatial context along the Z dimension will then not be taken into
+    account.
+    >>> config_2d = create_n2v_configuration(
+    ...     experiment_name="n2v_experiment",
+    ...     data_type="czi",
+    ...     axes="SCYX",
+    ...     patch_size=[64, 64],
+    ...     batch_size=32,
+    ...     num_epochs=100,
+    ...     n_channels=1,
+    ... )
+    >>> config_3d = create_n2v_configuration(
+    ...     experiment_name="n2v_experiment",
+    ...     data_type="czi",
+    ...     axes="SCZYX",
+    ...     patch_size=[16, 64, 64],
+    ...     batch_size=16,
+    ...     num_epochs=100,
+    ...     n_channels=1,
+    ... )
     """
     # if there are channels, we need to specify their number
     if "C" in axes and n_channels is None:
@@ -982,6 +1242,10 @@ def create_n2v_configuration(
         n_channels_out=n_channels,
         use_n2v2=use_n2v2,
         model_params=model_params,
+        optimizer=optimizer,
+        optimizer_params=optimizer_params,
+        lr_scheduler=lr_scheduler,
+        lr_scheduler_params=lr_scheduler_params,
     )
     algorithm_params["n2v_config"] = n2v_transform
 
@@ -1000,6 +1264,7 @@ def create_n2v_configuration(
     training_params = _create_training_configuration(
         num_epochs=num_epochs,
         logger=logger,
+        checkpoint_params=checkpoint_params,
     )
 
     return Configuration(

careamics/config/data/__init__.py

@@ -2,6 +2,8 @@
 
 __all__ = [
     "DataConfig",
+    "NGDataConfig",
 ]
 
 from .data_model import DataConfig
+from .ng_data_model import NGDataConfig

careamics/config/data/data_model.py

@@ -95,9 +95,9 @@ class DataConfig(BaseModel):
     )
 
     # Dataset configuration
-    data_type: Literal["array", "tiff", "custom"]
-    """Type of input data, numpy.ndarray (array) or paths (tiff and custom), as defined
-    in SupportedData."""
+    data_type: Literal["array", "tiff", "czi", "custom"]
+    """Type of input data, numpy.ndarray (array) or paths (tiff, czi, and custom), as
+    defined in SupportedData."""
 
     axes: str
     """Axes of the data, as defined in SupportedAxes."""