careamics 0.1.0rc6__py3-none-any.whl → 0.1.0rc8__py3-none-any.whl

careamics/careamist.py

@@ -1,9 +1,10 @@
 """A class to train, predict and export models in CAREamics."""
 
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union, overload
+from typing import Any, Callable, Literal, Optional, Union, overload
 
 import numpy as np
+from numpy.typing import NDArray
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import (
     Callback,
@@ -12,59 +13,64 @@ from pytorch_lightning.callbacks import (
 )
 from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger
 
-from careamics.callbacks import ProgressBarCallback
 from careamics.config import (
     Configuration,
-    create_inference_configuration,
     load_configuration,
 )
-from careamics.config.support import SupportedAlgorithm, SupportedData, SupportedLogger
+from careamics.config.support import (
+    SupportedAlgorithm,
+    SupportedArchitecture,
+    SupportedData,
+    SupportedLogger,
+)
 from careamics.dataset.dataset_utils import reshape_array
-from careamics.lightning_datamodule import CAREamicsTrainData
-from careamics.lightning_module import CAREamicsModule
-from careamics.lightning_prediction_datamodule import CAREamicsPredictData
-from careamics.lightning_prediction_loop import CAREamicsPredictionLoop
+from careamics.lightning import (
+    CAREamicsModule,
+    HyperParametersCallback,
+    PredictDataModule,
+    ProgressBarCallback,
+    TrainDataModule,
+    create_predict_datamodule,
+)
 from careamics.model_io import export_to_bmz, load_pretrained
-from careamics.transforms import Denormalize
+from careamics.prediction_utils import convert_outputs
 from careamics.utils import check_path_exists, get_logger
 
-from .callbacks import HyperParametersCallback
-
 logger = get_logger(__name__)
 
 LOGGER_TYPES = Optional[Union[TensorBoardLogger, WandbLogger]]
 
 
-# TODO napari callbacks
-# TODO: how to do AMP? How to continue training?
 class CAREamist:
     """Main CAREamics class, allowing training and prediction using various algorithms.
 
     Parameters
     ----------
-    source : Union[Path, str, Configuration]
+    source : pathlib.Path or str or CAREamics Configuration
         Path to a configuration file or a trained model.
-    work_dir : Optional[str], optional
+    work_dir : str, optional
         Path to working directory in which to save checkpoints and logs,
         by default None.
-    experiment_name : str, optional
-        Experiment name used for checkpoints, by default "CAREamics".
+    experiment_name : str, by default "CAREamics"
+        Experiment name used for checkpoints.
+    callbacks : list of Callback, optional
+        List of callbacks to use during training and prediction, by default None.
 
     Attributes
     ----------
-    model : CAREamicsKiln
+    model : CAREamicsModule
         CAREamics model.
     cfg : Configuration
         CAREamics configuration.
     trainer : Trainer
         PyTorch Lightning trainer.
-    experiment_logger : Optional[Union[TensorBoardLogger, WandbLogger]]
+    experiment_logger : TensorBoardLogger or WandbLogger
         Experiment logger, "wandb" or "tensorboard".
-    work_dir : Path
+    work_dir : pathlib.Path
         Working directory.
-    train_datamodule : Optional[CAREamicsWood]
+    train_datamodule : TrainDataModule
         Training datamodule.
-    pred_datamodule : Optional[CAREamicsClay]
+    pred_datamodule : PredictDataModule
         Prediction datamodule.
     """
 
@@ -74,6 +80,7 @@ class CAREamist:
         source: Union[Path, str],
         work_dir: Optional[str] = None,
         experiment_name: str = "CAREamics",
+        callbacks: Optional[list[Callback]] = None,
     ) -> None: ...
 
     @overload
@@ -82,6 +89,7 @@ class CAREamist:
         source: Configuration,
         work_dir: Optional[str] = None,
         experiment_name: str = "CAREamics",
+        callbacks: Optional[list[Callback]] = None,
     ) -> None: ...
 
     def __init__(
@@ -89,6 +97,7 @@ class CAREamist:
         source: Union[Path, str, Configuration],
         work_dir: Optional[Union[Path, str]] = None,
         experiment_name: str = "CAREamics",
+        callbacks: Optional[list[Callback]] = None,
     ) -> None:
         """
         Initialize CAREamist with a configuration object or a path.
@@ -105,13 +114,15 @@ class CAREamist:
 
         Parameters
         ----------
-        source : Union[Path, str, Configuration]
+        source : pathlib.Path or str or CAREamics Configuration
             Path to a configuration file or a trained model.
-        work_dir : Optional[str], optional
+        work_dir : str, optional
             Path to working directory in which to save checkpoints and logs,
             by default None.
         experiment_name : str, optional
             Experiment name used for checkpoints, by default "CAREamics".
+        callbacks : list of Callback, optional
+            List of callbacks to use during training and prediction, by default None.
 
         Raises
         ------
@@ -164,7 +175,7 @@ class CAREamist:
                 self.model, self.cfg = load_pretrained(source)
 
         # define the checkpoint saving callback
-        self.callbacks = self._define_callbacks()
+        self._define_callbacks(callbacks)
 
         # instantiate logger
         if self.cfg.training_config.has_logger():
@@ -188,32 +199,50 @@ class CAREamist:
             logger=self.experiment_logger,
         )
 
-        # change the prediction loop, necessary for tiled prediction
-        self.trainer.predict_loop = CAREamicsPredictionLoop(self.trainer)
-
         # place holder for the datamodules
-        self.train_datamodule: Optional[CAREamicsTrainData] = None
-        self.pred_datamodule: Optional[CAREamicsPredictData] = None
+        self.train_datamodule: Optional[TrainDataModule] = None
+        self.pred_datamodule: Optional[PredictDataModule] = None
 
-    def _define_callbacks(self) -> List[Callback]:
+    def _define_callbacks(self, callbacks: Optional[list[Callback]] = None) -> None:
         """
         Define the callbacks for the training loop.
 
-        Returns
-        -------
-        List[Callback]
-            List of callbacks to be used during training.
+        Parameters
+        ----------
+        callbacks : list of Callback, optional
+            List of callbacks to use during training and prediction, by default None.
         """
+        self.callbacks = [] if callbacks is None else callbacks
+
+        # check that user callbacks are not any of the CAREamics callbacks
+        for c in self.callbacks:
+            if isinstance(c, ModelCheckpoint) or isinstance(c, EarlyStopping):
+                raise ValueError(
+                    "ModelCheckpoint and EarlyStopping callbacks are already defined "
+                    "in CAREamics and should only be modified through the "
+                    "training configuration (see TrainingConfig)."
+                )
+
+            if isinstance(c, HyperParametersCallback) or isinstance(
+                c, ProgressBarCallback
+            ):
+                raise ValueError(
+                    "HyperParameter and ProgressBar callbacks are defined internally "
+                    "and should not be passed as callbacks."
+                )
+
         # checkpoint callback saves checkpoints during training
-        self.callbacks = [
-            HyperParametersCallback(self.cfg),
-            ModelCheckpoint(
-                dirpath=self.work_dir / Path("checkpoints"),
-                filename=self.cfg.experiment_name,
-                **self.cfg.training_config.checkpoint_callback.model_dump(),
-            ),
-            ProgressBarCallback(),
-        ]
+        self.callbacks.extend(
+            [
+                HyperParametersCallback(self.cfg),
+                ModelCheckpoint(
+                    dirpath=self.work_dir / Path("checkpoints"),
+                    filename=self.cfg.experiment_name,
+                    **self.cfg.training_config.checkpoint_callback.model_dump(),
+                ),
+                ProgressBarCallback(),
+            ]
+        )
 
         # early stopping callback
         if self.cfg.training_config.early_stopping_callback is not None:
@@ -221,16 +250,14 @@ class CAREamist:
                 EarlyStopping(self.cfg.training_config.early_stopping_callback)
             )
 
-        return self.callbacks
-
     def train(
         self,
         *,
-        datamodule: Optional[CAREamicsTrainData] = None,
-        train_source: Optional[Union[Path, str, np.ndarray]] = None,
-        val_source: Optional[Union[Path, str, np.ndarray]] = None,
-        train_target: Optional[Union[Path, str, np.ndarray]] = None,
-        val_target: Optional[Union[Path, str, np.ndarray]] = None,
+        datamodule: Optional[TrainDataModule] = None,
+        train_source: Optional[Union[Path, str, NDArray]] = None,
+        val_source: Optional[Union[Path, str, NDArray]] = None,
+        train_target: Optional[Union[Path, str, NDArray]] = None,
+        val_target: Optional[Union[Path, str, NDArray]] = None,
         use_in_memory: bool = True,
         val_percentage: float = 0.1,
         val_minimum_split: int = 1,
@@ -253,15 +280,15 @@ class CAREamist:
 
         Parameters
         ----------
-        datamodule : Optional[CAREamicsWood], optional
+        datamodule : TrainDataModule, optional
             Datamodule to train on, by default None.
-        train_source : Optional[Union[Path, str, np.ndarray]], optional
+        train_source : pathlib.Path or str or NDArray, optional
             Train source, if no datamodule is provided, by default None.
-        val_source : Optional[Union[Path, str, np.ndarray]], optional
+        val_source : pathlib.Path or str or NDArray, optional
             Validation source, if no datamodule is provided, by default None.
-        train_target : Optional[Union[Path, str, np.ndarray]], optional
+        train_target : pathlib.Path or str or NDArray, optional
             Train target source, if no datamodule is provided, by default None.
-        val_target : Optional[Union[Path, str, np.ndarray]], optional
+        val_target : pathlib.Path or str or NDArray, optional
             Validation target source, if no datamodule is provided, by default None.
         use_in_memory : bool, optional
             Use in memory dataset if possible, by default True.
@@ -355,17 +382,17 @@ class CAREamist:
 
             else:
                 raise ValueError(
-                    f"Invalid input, expected a str, Path, array or CAREamicsWood "
+                    f"Invalid input, expected a str, Path, array or TrainDataModule "
                     f"instance (got {type(train_source)})."
                 )
 
-    def _train_on_datamodule(self, datamodule: CAREamicsTrainData) -> None:
+    def _train_on_datamodule(self, datamodule: TrainDataModule) -> None:
         """
         Train the model on the provided datamodule.
 
         Parameters
         ----------
-        datamodule : CAREamicsWood
+        datamodule : TrainDataModule
             Datamodule to train on.
         """
         # record datamodule
@@ -375,10 +402,10 @@ class CAREamist:
 
     def _train_on_array(
         self,
-        train_data: np.ndarray,
-        val_data: Optional[np.ndarray] = None,
-        train_target: Optional[np.ndarray] = None,
-        val_target: Optional[np.ndarray] = None,
+        train_data: NDArray,
+        val_data: Optional[NDArray] = None,
+        train_target: Optional[NDArray] = None,
+        val_target: Optional[NDArray] = None,
         val_percentage: float = 0.1,
         val_minimum_split: int = 5,
     ) -> None:
@@ -387,13 +414,13 @@ class CAREamist:
 
         Parameters
         ----------
-        train_data : np.ndarray
+        train_data : NDArray
             Training data.
-        val_data : Optional[np.ndarray], optional
+        val_data : NDArray, optional
             Validation data, by default None.
-        train_target : Optional[np.ndarray], optional
+        train_target : NDArray, optional
             Train target data, by default None.
-        val_target : Optional[np.ndarray], optional
+        val_target : NDArray, optional
             Validation target data, by default None.
         val_percentage : float, optional
             Percentage of patches to use for validation, by default 0.1.
@@ -401,7 +428,7 @@ class CAREamist:
             Minimum number of patches to use for validation, by default 5.
         """
         # create datamodule
-        datamodule = CAREamicsTrainData(
+        datamodule = TrainDataModule(
             data_config=self.cfg.data_config,
             train_data=train_data,
             val_data=val_data,
@@ -429,13 +456,13 @@ class CAREamist:
 
         Parameters
         ----------
-        path_to_train_data : Union[Path, str]
+        path_to_train_data : pathlib.Path or str
             Path to the training data.
-        path_to_val_data : Optional[Union[Path, str]], optional
+        path_to_val_data : pathlib.Path or str, optional
             Path to validation data, by default None.
-        path_to_train_target : Optional[Union[Path, str]], optional
+        path_to_train_target : pathlib.Path or str, optional
             Path to train target data, by default None.
-        path_to_val_target : Optional[Union[Path, str]], optional
+        path_to_val_target : pathlib.Path or str, optional
             Path to validation target data, by default None.
         use_in_memory : bool, optional
             Use in memory dataset if possible, by default True.
@@ -457,7 +484,7 @@ class CAREamist:
             path_to_val_target = check_path_exists(path_to_val_target)
 
         # create datamodule
-        datamodule = CAREamicsTrainData(
+        datamodule = TrainDataModule(
             data_config=self.cfg.data_config,
             train_data=path_to_train_data,
             val_data=path_to_val_data,
@@ -473,11 +500,8 @@ class CAREamist:
 
     @overload
     def predict(  # numpydoc ignore=GL08
-        self,
-        source: CAREamicsPredictData,
-        *,
-        checkpoint: Optional[Literal["best", "last"]] = None,
-    ) -> Union[list, np.ndarray]: ...
+        self, source: PredictDataModule
+    ) -> Union[list[NDArray], NDArray]: ...
 
     @overload
     def predict(  # numpydoc ignore=GL08
@@ -485,59 +509,59 @@ class CAREamist:
         source: Union[Path, str],
         *,
         batch_size: int = 1,
-        tile_size: Optional[Tuple[int, ...]] = None,
-        tile_overlap: Tuple[int, ...] = (48, 48),
+        tile_size: Optional[tuple[int, ...]] = None,
+        tile_overlap: tuple[int, ...] = (48, 48),
         axes: Optional[str] = None,
         data_type: Optional[Literal["tiff", "custom"]] = None,
         tta_transforms: bool = True,
-        dataloader_params: Optional[Dict] = None,
+        dataloader_params: Optional[dict] = None,
         read_source_func: Optional[Callable] = None,
         extension_filter: str = "",
-        checkpoint: Optional[Literal["best", "last"]] = None,
-    ) -> Union[list, np.ndarray]: ...
+    ) -> Union[list[NDArray], NDArray]: ...
 
     @overload
     def predict(  # numpydoc ignore=GL08
         self,
-        source: np.ndarray,
+        source: NDArray,
         *,
         batch_size: int = 1,
-        tile_size: Optional[Tuple[int, ...]] = None,
-        tile_overlap: Tuple[int, ...] = (48, 48),
+        tile_size: Optional[tuple[int, ...]] = None,
+        tile_overlap: tuple[int, ...] = (48, 48),
         axes: Optional[str] = None,
         data_type: Optional[Literal["array"]] = None,
         tta_transforms: bool = True,
-        dataloader_params: Optional[Dict] = None,
-        checkpoint: Optional[Literal["best", "last"]] = None,
-    ) -> Union[list, np.ndarray]: ...
+        dataloader_params: Optional[dict] = None,
+    ) -> Union[list[NDArray], NDArray]: ...
 
     def predict(
         self,
-        source: Union[CAREamicsPredictData, Path, str, np.ndarray],
+        source: Union[PredictDataModule, Path, str, NDArray],
         *,
-        batch_size: int = 1,
-        tile_size: Optional[Tuple[int, ...]] = None,
-        tile_overlap: Tuple[int, ...] = (48, 48),
+        batch_size: Optional[int] = None,
+        tile_size: Optional[tuple[int, ...]] = None,
+        tile_overlap: tuple[int, ...] = (48, 48),
         axes: Optional[str] = None,
         data_type: Optional[Literal["array", "tiff", "custom"]] = None,
         tta_transforms: bool = True,
-        dataloader_params: Optional[Dict] = None,
+        dataloader_params: Optional[dict] = None,
         read_source_func: Optional[Callable] = None,
         extension_filter: str = "",
-        checkpoint: Optional[Literal["best", "last"]] = None,
         **kwargs: Any,
-    ) -> Union[List[np.ndarray], np.ndarray]:
+    ) -> Union[list[NDArray], NDArray]:
         """
         Make predictions on the provided data.
 
-        Input can be a CAREamicsClay instance, a path to a data file, or a numpy array.
+        Input can be a CAREamicsPredData instance, a path to a data file, or a numpy
+        array.
 
         If `data_type`, `axes` and `tile_size` are not provided, the training
         configuration parameters will be used, with the `patch_size` instead of
         `tile_size`.
 
         Test-time augmentation (TTA) can be switched off using the `tta_transforms`
-        parameter.
+        parameter. The TTA augmentation applies all possible flip and 90 degrees
+        rotations to the prediction input and averages the predictions. TTA augmentation
+        should not be used if you did not train with these augmentations.
 
         Note that if you are using a UNet model and tiling, the tile size must be
         divisible in every dimension by 2**d, where d is the depth of the model. This
@@ -547,242 +571,136 @@ class CAREamist:
 
         Parameters
         ----------
-        source : Union[CAREamicsClay, Path, str, np.ndarray]
+        source : CAREamicsPredData, pathlib.Path, str or numpy.ndarray
             Data to predict on.
-        batch_size : int, optional
-            Batch size for prediction, by default 1.
-        tile_size : Optional[Tuple[int, ...]], optional
-            Size of the tiles to use for prediction, by default None.
-        tile_overlap : Tuple[int, ...], optional
-            Overlap between tiles, by default (48, 48).
-        axes : Optional[str], optional
+        batch_size : int, default=1
+            Batch size for prediction.
+        tile_size : tuple of int, optional
+            Size of the tiles to use for prediction.
+        tile_overlap : tuple of int, default=(48, 48)
+            Overlap between tiles.
+        axes : str, optional
             Axes of the input data, by default None.
-        data_type : Optional[Literal["array", "tiff", "custom"]], optional
-            Type of the input data, by default None.
-        tta_transforms : bool, optional
-            Whether to apply test-time augmentation, by default True.
-        dataloader_params : Optional[Dict], optional
-            Parameters to pass to the dataloader, by default None.
-        read_source_func : Optional[Callable], optional
-            Function to read the source data, by default None.
-        extension_filter : str, optional
-            Filter for the file extension, by default "".
-        checkpoint : Optional[Literal["best", "last"]], optional
-            Checkpoint to use for prediction, by default None.
+        data_type : {"array", "tiff", "custom"}, optional
+            Type of the input data.
+        tta_transforms : bool, default=True
+            Whether to apply test-time augmentation.
+        dataloader_params : dict, optional
+            Parameters to pass to the dataloader.
+        read_source_func : Callable, optional
+            Function to read the source data.
+        extension_filter : str, default=""
+            Filter for the file extension.
         **kwargs : Any
             Unused.
 
         Returns
         -------
-        Union[List[np.ndarray], np.ndarray]
+        list of NDArray or NDArray
             Predictions made by the model.
 
-        Raises
-        ------
-        ValueError
-            If the input is not a CAREamicsClay instance, a path or a numpy array.
-        """
-        if isinstance(source, CAREamicsPredictData):
-            # record datamodule
-            self.pred_datamodule = source
-
-            return self.trainer.predict(
-                model=self.model, datamodule=source, ckpt_path=checkpoint
-            )
-        else:
-            if self.cfg is None:
-                raise ValueError(
-                    "No configuration found. Train a model or load from a "
-                    "checkpoint before predicting."
-                )
-            # create predict config, reuse training config if parameters missing
-            prediction_config = create_inference_configuration(
-                configuration=self.cfg,
-                tile_size=tile_size,
-                tile_overlap=tile_overlap,
-                data_type=data_type,
-                axes=axes,
-                tta_transforms=tta_transforms,
-                batch_size=batch_size,
-            )
-
-            # remove batch from dataloader parameters (priority given to config)
-            if dataloader_params is None:
-                dataloader_params = {}
-            if "batch_size" in dataloader_params:
-                del dataloader_params["batch_size"]
-
-            if isinstance(source, Path) or isinstance(source, str):
-                # Check the source
-                source_path = check_path_exists(source)
-
-                # create datamodule
-                datamodule = CAREamicsPredictData(
-                    pred_config=prediction_config,
-                    pred_data=source_path,
-                    read_source_func=read_source_func,
-                    extension_filter=extension_filter,
-                    dataloader_params=dataloader_params,
-                )
-
-                # record datamodule
-                self.pred_datamodule = datamodule
-
-                return self.trainer.predict(
-                    model=self.model, datamodule=datamodule, ckpt_path=checkpoint
-                )
-
-            elif isinstance(source, np.ndarray):
-                # create datamodule
-                datamodule = CAREamicsPredictData(
-                    pred_config=prediction_config,
-                    pred_data=source,
-                    dataloader_params=dataloader_params,
-                )
-
-                # record datamodule
-                self.pred_datamodule = datamodule
-
-                return self.trainer.predict(
-                    model=self.model, datamodule=datamodule, ckpt_path=checkpoint
-                )
-
-            else:
-                raise ValueError(
-                    f"Invalid input. Expected a CAREamicsWood instance, paths or "
-                    f"np.ndarray (got {type(source)})."
-                )
-
-    def _create_data_for_bmz(
-        self,
-        input_array: Optional[np.ndarray] = None,
-    ) -> np.ndarray:
-        """Create data for BMZ export.
-
-        If no `input_array` is provided, this method checks if there is a prediction
-        datamodule, or a training data module, to extract a patch. If none exists,
-        then a random aray is created.
-
-        If there is a non-singleton batch dimension, this method returns only the first
-        element.
-
-        Parameters
-        ----------
-        input_array : Optional[np.ndarray], optional
-            Input array, by default None.
-
-        Returns
-        -------
-        np.ndarray
-            Input data for BMZ export.
-
         Raises
         ------
         ValueError
             If mean and std are not provided in the configuration.
+        ValueError
+            If tile size is not divisible by 2**depth for UNet models.
+        ValueError
+            If tile overlap is not specified.
         """
-        if input_array is None:
-            if self.cfg.data_config.mean is None or self.cfg.data_config.std is None:
-                raise ValueError(
-                    "Mean and std cannot be None in the configuration in order to"
-                    "export to the BMZ format. Was the model trained?"
-                )
-
-            # generate images, priority is given to the prediction data module
-            if self.pred_datamodule is not None:
-                # unpack a batch, ignore masks or targets
-                input_patch, *_ = next(iter(self.pred_datamodule.predict_dataloader()))
-
-                # convert torch.Tensor to numpy
-                input_patch = input_patch.numpy()
-
-                # denormalize
-                denormalize = Denormalize(
-                    mean=self.cfg.data_config.mean, std=self.cfg.data_config.std
-                )
-                input_patch, _ = denormalize(input_patch)
-
-            elif self.train_datamodule is not None:
-                input_patch, *_ = next(iter(self.train_datamodule.train_dataloader()))
-                input_patch = input_patch.numpy()
-
-                # denormalize
-                denormalize = Denormalize(
-                    mean=self.cfg.data_config.mean, std=self.cfg.data_config.std
-                )
-                input_patch, _ = denormalize(input_patch)
-            else:
-                # create a random input array
-                input_patch = np.random.normal(
-                    loc=self.cfg.data_config.mean,
-                    scale=self.cfg.data_config.std,
-                    size=self.cfg.data_config.patch_size,
-                ).astype(np.float32)[
-                    np.newaxis, np.newaxis, ...
-                ]  # add S & C dimensions
-        else:
-            # potentially correct shape
-            input_patch = reshape_array(input_array, self.cfg.data_config.axes)
-
-        # if this a batch
-        if input_patch.shape[0] > 1:
-            input_patch = input_patch[[0], ...]  # keep singleton dim
+        if (
+            self.cfg.data_config.image_means is None
+            or self.cfg.data_config.image_stds is None
+        ):
+            raise ValueError("Mean and std must be provided in the configuration.")
+
+        # tile size for UNets
+        if tile_size is not None:
+            model = self.cfg.algorithm_config.model
+
+            if model.architecture == SupportedArchitecture.UNET.value:
+                # tile size must be equal to k*2^n, where n is the number of pooling
+                # layers (equal to the depth) and k is an integer
+                depth = model.depth
+                tile_increment = 2**depth
+
+                for i, t in enumerate(tile_size):
+                    if t % tile_increment != 0:
+                        raise ValueError(
+                            f"Tile size must be divisible by {tile_increment} along "
+                            f"all axes (got {t} for axis {i}). If your image size is "
+                            f"smaller along one axis (e.g. Z), consider padding the "
+                            f"image."
+                        )
+
+            # tile overlaps must be specified
+            if tile_overlap is None:
+                raise ValueError("Tile overlap must be specified.")
+
+        # create the prediction
+        self.pred_datamodule = create_predict_datamodule(
+            pred_data=source,
+            data_type=data_type or self.cfg.data_config.data_type,
+            axes=axes or self.cfg.data_config.axes,
+            image_means=self.cfg.data_config.image_means,
+            image_stds=self.cfg.data_config.image_stds,
+            tile_size=tile_size,
+            tile_overlap=tile_overlap,
+            batch_size=batch_size or self.cfg.data_config.batch_size,
+            tta_transforms=tta_transforms,
+            read_source_func=read_source_func,
+            extension_filter=extension_filter,
+            dataloader_params=dataloader_params,
+        )
 
-        return input_patch
+        # predict
+        predictions = self.trainer.predict(
+            model=self.model, datamodule=self.pred_datamodule
+        )
+        return convert_outputs(predictions, self.pred_datamodule.tiled)
 
     def export_to_bmz(
         self,
         path: Union[Path, str],
         name: str,
-        authors: List[dict],
-        input_array: Optional[np.ndarray] = None,
+        input_array: NDArray,
+        authors: list[dict],
         general_description: str = "",
-        channel_names: Optional[List[str]] = None,
+        channel_names: Optional[list[str]] = None,
         data_description: Optional[str] = None,
     ) -> None:
         """Export the model to the BioImage Model Zoo format.
 
-        Input array must be of shape SC(Z)YX, with S and C singleton dimensions.
+        Input array must be of the same dimensions as the axes recorded in the
+        configuration of the `CAREamist`.
 
         Parameters
         ----------
-        path : Union[Path, str]
+        path : pathlib.Path or str
             Path to save the model.
         name : str
             Name of the model.
-        authors : List[dict]
+        input_array : NDArray
+            Input array used to validate the model and as example.
+        authors : list of dict
             List of authors of the model.
-        input_array : Optional[np.ndarray], optional
-            Input array for the model, must be of shape SC(Z)YX, by default None.
         general_description : str
             General description of the model, used in the metadata of the BMZ archive.
-        channel_names : Optional[List[str]], optional
+        channel_names : list of str, optional
             Channel names, by default None.
-        data_description : Optional[str], optional
+        data_description : str, optional
             Description of the data, by default None.
         """
-        input_patch = self._create_data_for_bmz(input_array)
+        # TODO: add in docs that it is expected that input_array dimensions match
+        # those in data_config
 
-        # axes need to be reformated for the export because reshaping was done in the
-        # datamodule
-        if "Z" in self.cfg.data_config.axes:
-            axes = "SCZYX"
-        else:
-            axes = "SCYX"
-
-        # predict output, remove extra dimensions for the purpose of the prediction
         output_patch = self.predict(
-            input_patch,
+            input_array,
             data_type=SupportedData.ARRAY.value,
-            axes=axes,
             tta_transforms=False,
         )
-
-        if not isinstance(output_patch, np.ndarray):
-            raise ValueError(
-                f"Numpy array required for export to BioImage Model Zoo, got "
-                f"{type(output_patch)}."
-            )
+        output = np.concatenate(output_patch, axis=0)
+        input_array = reshape_array(input_array, self.cfg.data_config.axes)
 
         export_to_bmz(
             model=self.model,
@@ -791,8 +709,8 @@ class CAREamist:
             name=name,
             general_description=general_description,
             authors=authors,
-            input_array=input_patch,
-            output_array=output_patch,
+            input_array=input_array,
+            output_array=output,
             channel_names=channel_names,
             data_description=data_description,
         )