careamics 0.1.0rc5__py3-none-any.whl → 0.1.0rc7__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,
@@ -15,55 +16,54 @@ 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.inference_model import TRANSFORMS_UNION
 from careamics.config.support import SupportedAlgorithm, 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.model_io import export_to_bmz, load_pretrained
+from careamics.prediction_utils import convert_outputs, create_pred_datamodule
 from careamics.utils import check_path_exists, get_logger
 
 from .callbacks import HyperParametersCallback
+from .lightning_prediction_datamodule import CAREamicsPredictData
 
 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 : CAREamicsTrainData
         Training datamodule.
-    pred_datamodule : Optional[CAREamicsClay]
+    pred_datamodule : CAREamicsPredictData
         Prediction datamodule.
     """
 
@@ -73,6 +73,7 @@ class CAREamist:
         source: Union[Path, str],
         work_dir: Optional[str] = None,
         experiment_name: str = "CAREamics",
+        callbacks: Optional[list[Callback]] = None,
     ) -> None: ...
 
     @overload
@@ -81,6 +82,7 @@ class CAREamist:
         source: Configuration,
         work_dir: Optional[str] = None,
         experiment_name: str = "CAREamics",
+        callbacks: Optional[list[Callback]] = None,
     ) -> None: ...
 
     def __init__(
@@ -88,6 +90,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.
@@ -104,13 +107,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
         ------
@@ -163,7 +168,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():
@@ -187,32 +192,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
 
-    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:
@@ -220,16 +243,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,
+        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,
@@ -252,15 +273,15 @@ class CAREamist:
 
         Parameters
         ----------
-        datamodule : Optional[CAREamicsWood], optional
+        datamodule : CAREamicsTrainData, 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.
@@ -354,7 +375,7 @@ class CAREamist:
 
             else:
                 raise ValueError(
-                    f"Invalid input, expected a str, Path, array or CAREamicsWood "
+                    f"Invalid input, expected a str, Path, array or CAREamicsTrainData "
                     f"instance (got {type(train_source)})."
                 )
 
@@ -364,7 +385,7 @@ class CAREamist:
 
         Parameters
         ----------
-        datamodule : CAREamicsWood
+        datamodule : CAREamicsTrainData
             Datamodule to train on.
         """
         # record datamodule
@@ -374,10 +395,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:
@@ -386,13 +407,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.
@@ -428,13 +449,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.
@@ -476,7 +497,7 @@ class CAREamist:
         source: CAREamicsPredictData,
         *,
         checkpoint: Optional[Literal["best", "last"]] = None,
-    ) -> Union[list, np.ndarray]: ...
+    ) -> Union[list[NDArray], NDArray]: ...
 
     @overload
     def predict(  # numpydoc ignore=GL08
@@ -484,64 +505,62 @@ 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,
-        transforms: Optional[List[TRANSFORMS_UNION]] = 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,
-        transforms: Optional[List[TRANSFORMS_UNION]] = None,
         tta_transforms: bool = True,
-        dataloader_params: Optional[Dict] = None,
+        dataloader_params: Optional[dict] = None,
         checkpoint: Optional[Literal["best", "last"]] = None,
-    ) -> Union[list, np.ndarray]: ...
+    ) -> Union[list[NDArray], NDArray]: ...
 
     def predict(
         self,
-        source: Union[CAREamicsPredictData, Path, str, np.ndarray],
+        source: Union[CAREamicsPredictData, 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,
-        transforms: Optional[List[TRANSFORMS_UNION]] = 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`.
 
-        The default transforms are defined in the `InferenceModel` Pydantic model.
-
         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
@@ -551,181 +570,96 @@ 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.
-        transforms : Optional[List[TRANSFORMS_UNION]], optional
-            List of transforms to apply to the 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.
+        checkpoint : {"best", "last"}, optional
+            Checkpoint to use for prediction.
         **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,
-                transforms=transforms,
-                tta_transforms=tta_transforms,
-                batch_size=batch_size,
+        # Reuse batch size if not provided explicitly
+        if batch_size is None:
+            batch_size = (
+                self.train_datamodule.batch_size
+                if self.train_datamodule
+                else self.cfg.data_config.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
-                )
+        self.pred_datamodule = create_pred_datamodule(
+            source=source,
+            config=self.cfg,
+            batch_size=batch_size,
+            tile_size=tile_size,
+            tile_overlap=tile_overlap,
+            axes=axes,
+            data_type=data_type,
+            tta_transforms=tta_transforms,
+            dataloader_params=dataloader_params,
+            read_source_func=read_source_func,
+            extension_filter=extension_filter,
+        )
 
-            else:
-                raise ValueError(
-                    f"Invalid input. Expected a CAREamicsWood instance, paths or "
-                    f"np.ndarray (got {type(source)})."
-                )
+        predictions = self.trainer.predict(
+            model=self.model, datamodule=self.pred_datamodule, ckpt_path=checkpoint
+        )
+        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.
         """
-        if input_array is None:
-            # 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()
-            elif self.train_datamodule is not None:
-                input_patch, *_ = next(iter(self.train_datamodule.train_dataloader()))
-                input_patch = input_patch.numpy()
-            else:
-                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?"
-                    )
-
-                # 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:
-            input_patch = input_array
-
-        # if there is a batch dimension
-        if input_patch.shape[0] > 1:
-            input_patch = input_patch[0:1, ...]  # keep singleton dim
+        input_patch = reshape_array(input_array, self.cfg.data_config.axes)
 
         # axes need to be reformated for the export because reshaping was done in the
         # datamodule
@@ -742,11 +676,10 @@ class CAREamist:
             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)}."
-            )
+        if isinstance(output_patch, list):
+            output = np.concatenate(output_patch, axis=0)
+        else:
+            output = output_patch
 
         export_to_bmz(
             model=self.model,
@@ -756,7 +689,7 @@ class CAREamist:
             general_description=general_description,
             authors=authors,
             input_array=input_patch,
-            output_array=output_patch,
+            output_array=output,
             channel_names=channel_names,
             data_description=data_description,
         )