careamics 0.1.0rc3__py3-none-any.whl → 0.1.0rc5__py3-none-any.whl

careamics/lightning_datamodule.py

@@ -1,12 +1,13 @@
+"""Training and validation Lightning data modules."""
+
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Literal, Optional, Union
 
 import numpy as np
 import pytorch_lightning as L
-from albumentations import Compose
 from torch.utils.data import DataLoader
 
-from careamics.config import DataModel
+from careamics.config import DataConfig
 from careamics.config.data_model import TRANSFORMS_UNION
 from careamics.config.support import SupportedData
 from careamics.dataset.dataset_utils import (
@@ -28,9 +29,9 @@ DatasetType = Union[InMemoryDataset, PathIterableDataset]
 logger = get_logger(__name__)
 
 
-class CAREamicsWood(L.LightningDataModule):
+class CAREamicsTrainData(L.LightningDataModule):
     """
-    LightningDataModule for training and validation datasets.
+    CAREamics Ligthning training and validation data module.
 
     The data module can be used with Path, str or numpy arrays. In the case of
     numpy arrays, it loads and computes all the patches in memory. For Path and str
@@ -53,11 +54,70 @@ class CAREamicsWood(L.LightningDataModule):
 
     You can also provide a `fnmatch` and `Path.rglob` compatible expression (e.g.
     "*.czi") to filter the files extension using `extension_filter`.
+
+    Parameters
+    ----------
+    data_config : DataModel
+        Pydantic model for CAREamics data configuration.
+    train_data : Union[Path, str, np.ndarray]
+        Training data, can be a path to a folder, a file or a numpy array.
+    val_data : Optional[Union[Path, str, np.ndarray]], optional
+        Validation data, can be a path to a folder, a file or a numpy array, by
+        default None.
+    train_data_target : Optional[Union[Path, str, np.ndarray]], optional
+        Training target data, can be a path to a folder, a file or a numpy array, by
+        default None.
+    val_data_target : Optional[Union[Path, str, np.ndarray]], optional
+        Validation target data, can be a path to a folder, a file or a numpy array,
+        by default None.
+    read_source_func : Optional[Callable], optional
+        Function to read the source data, by default None. Only used for `custom`
+        data type (see DataModel).
+    extension_filter : str, optional
+        Filter for file extensions, by default "". Only used for `custom` data types
+        (see DataModel).
+    val_percentage : float, optional
+        Percentage of the training data to use for validation, by default 0.1. Only
+        used if `val_data` is None.
+    val_minimum_split : int, optional
+        Minimum number of patches or files to split from the training data for
+        validation, by default 5. Only used if `val_data` is None.
+    use_in_memory : bool, optional
+        Use in memory dataset if possible, by default True.
+
+    Attributes
+    ----------
+    data_config : DataModel
+        CAREamics data configuration.
+    data_type : SupportedData
+        Expected data type, one of "tiff", "array" or "custom".
+    batch_size : int
+        Batch size.
+    use_in_memory : bool
+        Whether to use in memory dataset if possible.
+    train_data : Union[Path, str, np.ndarray]
+        Training data.
+    val_data : Optional[Union[Path, str, np.ndarray]]
+        Validation data.
+    train_data_target : Optional[Union[Path, str, np.ndarray]]
+        Training target data.
+    val_data_target : Optional[Union[Path, str, np.ndarray]]
+        Validation target data.
+    val_percentage : float
+        Percentage of the training data to use for validation, if no validation data is
+        provided.
+    val_minimum_split : int
+        Minimum number of patches or files to split from the training data for
+        validation, if no validation data is provided.
+    read_source_func : Optional[Callable]
+        Function to read the source data, used if `data_type` is `custom`.
+    extension_filter : str
+        Filter for file extensions, used if `data_type` is `custom`.
     """
 
     def __init__(
         self,
-        data_config: DataModel,
+        data_config: DataConfig,
         train_data: Union[Path, str, np.ndarray],
         val_data: Optional[Union[Path, str, np.ndarray]] = None,
         train_data_target: Optional[Union[Path, str, np.ndarray]] = None,
@@ -98,6 +158,8 @@ class CAREamicsWood(L.LightningDataModule):
         val_minimum_split : int, optional
             Minimum number of patches or files to split from the training data for
             validation, by default 5. Only used if `val_data` is None.
+        use_in_memory : bool, optional
+            Use in memory dataset if possible, by default True.
 
         Raises
         ------
@@ -128,25 +190,30 @@ class CAREamicsWood(L.LightningDataModule):
         if data_config.data_type == SupportedData.CUSTOM and read_source_func is None:
             raise ValueError(
                 f"Data type {SupportedData.CUSTOM} is not allowed without "
-                f"specifying a `read_source_func`."
+                f"specifying a `read_source_func` and an `extension_filer`."
             )
 
-        # and that arrays are passed, if array type specified
-        elif data_config.data_type == SupportedData.ARRAY and not isinstance(
-            train_data, np.ndarray
+        # check correct input type
+        if (
+            isinstance(train_data, np.ndarray)
+            and data_config.data_type != SupportedData.ARRAY
         ):
             raise ValueError(
-                f"Expected array input (see configuration.data.data_type), but got "
-                f"{type(train_data)} instead."
+                f"Received a numpy array as input, but the data type was set to "
+                f"{data_config.data_type}. Set the data type in the configuration "
+                f"to {SupportedData.ARRAY} to train on numpy arrays."
             )
 
         # and that Path or str are passed, if tiff file type specified
-        elif data_config.data_type == SupportedData.TIFF and (
-            not isinstance(train_data, Path) and not isinstance(train_data, str)
+        elif (isinstance(train_data, Path) or isinstance(train_data, str)) and (
+            data_config.data_type != SupportedData.TIFF
+            and data_config.data_type != SupportedData.CUSTOM
         ):
             raise ValueError(
-                f"Expected Path or str input (see configuration.data.data_type), "
-                f"but got {type(train_data)} instead."
+                f"Received a path as input, but the data type was neither set to "
+                f"{SupportedData.TIFF} nor {SupportedData.CUSTOM}. Set the data type "
+                f"in the configuration to {SupportedData.TIFF} or "
+                f"{SupportedData.CUSTOM} to train on files."
             )
 
         # configuration
@@ -231,7 +298,15 @@ class CAREamicsWood(L.LightningDataModule):
                 validate_source_target_files(self.val_files, self.val_target_files)
 
     def setup(self, *args: Any, **kwargs: Any) -> None:
-        """Hook called at the beginning of fit, validate, or predict."""
+        """Hook called at the beginning of fit, validate, or predict.
+
+        Parameters
+        ----------
+        *args : Any
+            Unused.
+        **kwargs : Any
+            Unused.
+        """
         # if numpy array
         if self.data_type == SupportedData.ARRAY:
             # train dataset
@@ -266,9 +341,9 @@ class CAREamicsWood(L.LightningDataModule):
                 self.train_dataset = InMemoryDataset(
                     data_config=self.data_config,
                     inputs=self.train_files,
-                    data_target=self.train_target_files
-                    if self.train_data_target
-                    else None,
+                    data_target=(
+                        self.train_target_files if self.train_data_target else None
+                    ),
                     read_source_func=self.read_source_func,
                 )
 
@@ -277,9 +352,9 @@ class CAREamicsWood(L.LightningDataModule):
                     self.val_dataset = InMemoryDataset(
                         data_config=self.data_config,
                         inputs=self.val_files,
-                        data_target=self.val_target_files
-                        if self.val_data_target
-                        else None,
+                        data_target=(
+                            self.val_target_files if self.val_data_target else None
+                        ),
                         read_source_func=self.read_source_func,
                     )
                 else:
@@ -295,9 +370,9 @@ class CAREamicsWood(L.LightningDataModule):
                 self.train_dataset = PathIterableDataset(
                     data_config=self.data_config,
                     src_files=self.train_files,
-                    target_files=self.train_target_files
-                    if self.train_data_target
-                    else None,
+                    target_files=(
+                        self.train_target_files if self.train_data_target else None
+                    ),
                     read_source_func=self.read_source_func,
                 )
 
@@ -307,9 +382,9 @@ class CAREamicsWood(L.LightningDataModule):
                     self.val_dataset = PathIterableDataset(
                         data_config=self.data_config,
                         src_files=self.val_files,
-                        target_files=self.val_target_files
-                        if self.val_data_target
-                        else None,
+                        target_files=(
+                            self.val_target_files if self.val_data_target else None
+                        ),
                         read_source_func=self.read_source_func,
                     )
                 elif len(self.train_files) <= self.val_minimum_split:
@@ -353,9 +428,12 @@ class CAREamicsWood(L.LightningDataModule):
         )
 
 
-class CAREamicsTrainDataModule(CAREamicsWood):
+class TrainingDataWrapper(CAREamicsTrainData):
     """
-    LightningDataModule wrapper for training and validation datasets.
+    Wrapper around the CAREamics Lightning training data module.
+
+    This class is used to explicitely pass the parameters usually contained in a
+    `data_model` configuration.
 
     Since the lightning datamodule has no access to the model, make sure that the
     parameters passed to the datamodule are consistent with the model's requirements and
@@ -374,8 +452,7 @@ class CAREamicsTrainDataModule(CAREamicsWood):
     In particular, N2V requires a specific transformation (N2V manipulates), which is
     not compatible with supervised training. The default transformations applied to the
     training patches are defined in `careamics.config.data_model`. To use different
-    transformations, pass a list of transforms or an albumentation `Compose` as
-    `transforms` parameter. See examples for more details.
+    transformations, pass a list of transforms. See examples for more details.
 
     By default, CAREamics only supports types defined in
     `careamics.config.support.SupportedData`. To read custom data types, you can set
@@ -410,7 +487,7 @@ class CAREamicsTrainDataModule(CAREamicsWood):
         Batch size.
     val_data : Optional[Union[str, Path]], optional
         Validation data, by default None.
-    transforms : Optional[Union[List[TRANSFORMS_UNION], Compose]], optional
+    transforms : List[TRANSFORMS_UNION], optional
         List of transforms to apply to training patches. If None, default transforms
         are applied.
     train_target_data : Optional[Union[str, Path]], optional
@@ -442,11 +519,11 @@ class CAREamicsTrainDataModule(CAREamicsWood):
 
     Examples
     --------
-    Create a CAREamicsTrainDataModule with default transforms with a numpy array:
+    Create a TrainingDataWrapper with default transforms with a numpy array:
     >>> import numpy as np
-    >>> from careamics import CAREamicsTrainDataModule
+    >>> from careamics import TrainingDataWrapper
     >>> my_array = np.arange(256).reshape(16, 16)
-    >>> data_module = CAREamicsTrainDataModule(
+    >>> data_module = TrainingDataWrapper(
     ...     train_data=my_array,
     ...     data_type="array",
     ...     patch_size=(8, 8),
@@ -457,12 +534,12 @@ class CAREamicsTrainDataModule(CAREamicsWood):
     For custom data types (those not supported by CAREamics), then one can pass a read
     function and a filter for the files extension:
     >>> import numpy as np
-    >>> from careamics import CAREamicsTrainDataModule
+    >>> from careamics import TrainingDataWrapper
     >>>
     >>> def read_npy(path):
     ...     return np.load(path)
     >>>
-    >>> data_module = CAREamicsTrainDataModule(
+    >>> data_module = TrainingDataWrapper(
     ...     train_data="path/to/data",
     ...     data_type="custom",
     ...     patch_size=(8, 8),
@@ -475,7 +552,7 @@ class CAREamicsTrainDataModule(CAREamicsWood):
     If you want to use a different set of transformations, you can pass a list of
     transforms:
     >>> import numpy as np
-    >>> from careamics import CAREamicsTrainDataModule
+    >>> from careamics import TrainingDataWrapper
     >>> from careamics.config.support import SupportedTransform
     >>> my_array = np.arange(256).reshape(16, 16)
     >>> my_transforms = [
@@ -488,7 +565,7 @@ class CAREamicsTrainDataModule(CAREamicsWood):
     ...         "name": SupportedTransform.N2V_MANIPULATE.value,
     ...     }
     ... ]
-    >>> data_module = CAREamicsTrainDataModule(
+    >>> data_module = TrainingDataWrapper(
     ...     train_data=my_array,
     ...     data_type="array",
     ...     patch_size=(8, 8),
@@ -506,7 +583,7 @@ class CAREamicsTrainDataModule(CAREamicsWood):
         axes: str,
         batch_size: int,
         val_data: Optional[Union[str, Path]] = None,
-        transforms: Optional[Union[List[TRANSFORMS_UNION], Compose]] = None,
+        transforms: Optional[List[TRANSFORMS_UNION]] = None,
         train_target_data: Optional[Union[str, Path]] = None,
         val_target_data: Optional[Union[str, Path]] = None,
         read_source_func: Optional[Callable] = None,
@@ -539,8 +616,8 @@ class CAREamicsTrainDataModule(CAREamicsWood):
         In particular, N2V requires a specific transformation (N2V manipulates), which
         is not compatible with supervised training. The default transformations applied
         to the training patches are defined in `careamics.config.data_model`. To use
-        different transformations, pass a list of transforms or an albumentation
-        `Compose` as `transforms` parameter. See examples for more details.
+        different transformations, pass a list of transforms. See examples for more
+        details.
 
         By default, CAREamics only supports types defined in
         `careamics.config.support.SupportedData`. To read custom data types, you can set
@@ -577,7 +654,7 @@ class CAREamicsTrainDataModule(CAREamicsWood):
             Batch size.
         val_data : Optional[Union[str, Path]], optional
             Validation data, by default None.
-        transforms : Optional[Union[List[TRANSFORMS_UNION], Compose]], optional
+        transforms : Optional[List[TRANSFORMS_UNION]], optional
             List of transforms to apply to training patches. If None, default transforms
             are applied.
         train_target_data : Optional[Union[str, Path]], optional
@@ -628,13 +705,10 @@ class CAREamicsTrainDataModule(CAREamicsWood):
             data_dict["transforms"] = transforms
 
         # validate configuration
-        self.data_config = DataModel(**data_dict)
+        self.data_config = DataConfig(**data_dict)
 
         # N2V specific checks, N2V, structN2V, and transforms
-        if (
-            self.data_config.has_transform_list()
-            and self.data_config.has_n2v_manipulate()
-        ):
+        if self.data_config.has_n2v_manipulate():
             # there is not target, n2v2 and structN2V can be changed
             if train_target_data is None:
                 self.data_config.set_N2V2(use_n2v2)

careamics/lightning_module.py

@@ -3,7 +3,7 @@ from typing import Any, Optional, Union
 import pytorch_lightning as L
 from torch import Tensor, nn
 
-from careamics.config import AlgorithmModel
+from careamics.config import AlgorithmConfig
 from careamics.config.support import (
     SupportedAlgorithm,
     SupportedArchitecture,
@@ -17,7 +17,7 @@ from careamics.transforms import Denormalize, ImageRestorationTTA
 from careamics.utils.torch_utils import get_optimizer, get_scheduler
 
 
-class CAREamicsKiln(L.LightningModule):
+class CAREamicsModule(L.LightningModule):
     """
     CAREamics Lightning module.
 
@@ -38,7 +38,7 @@ class CAREamicsKiln(L.LightningModule):
         Learning rate scheduler name.
     """
 
-    def __init__(self, algorithm_config: Union[AlgorithmModel, dict]) -> None:
+    def __init__(self, algorithm_config: Union[AlgorithmConfig, dict]) -> None:
         """
         CAREamics Lightning module.
 
@@ -53,7 +53,7 @@ class CAREamicsKiln(L.LightningModule):
         super().__init__()
         # if loading from a checkpoint, AlgorithmModel needs to be instantiated
         if isinstance(algorithm_config, dict):
-            algorithm_config = AlgorithmModel(**algorithm_config)
+            algorithm_config = AlgorithmConfig(**algorithm_config)
 
         # create model and loss function
         self.model: nn.Module = model_factory(algorithm_config.model)
@@ -162,7 +162,7 @@ class CAREamicsKiln(L.LightningModule):
             mean=self._trainer.datamodule.predict_dataset.mean,
             std=self._trainer.datamodule.predict_dataset.std,
         )
-        denormalized_output = denorm(image=output)["image"]
+        denormalized_output, _ = denorm(patch=output)
 
         if len(aux) > 0:
             return denormalized_output, aux
@@ -192,7 +192,7 @@ class CAREamicsKiln(L.LightningModule):
         }
 
 
-class CAREamicsModule(CAREamicsKiln):
+class CAREamicsModuleWrapper(CAREamicsModule):
     """Class defining the API for CAREamics Lightning layer.
 
     This class exposes parameters used to create an AlgorithmModel instance, triggering
@@ -287,6 +287,6 @@ class CAREamicsModule(CAREamicsKiln):
         algorithm_configuration["model"] = model_configuration
 
         # call the parent init using an AlgorithmModel instance
-        super().__init__(AlgorithmModel(**algorithm_configuration))
+        super().__init__(AlgorithmConfig(**algorithm_configuration))
 
         # TODO add load_from_checkpoint wrapper

careamics/lightning_prediction_datamodule.py

@@ -1,13 +1,14 @@
+"""Prediction Lightning data modules."""
+
 from pathlib import Path
 from typing import Any, Callable, List, Literal, Optional, Tuple, Union
 
 import numpy as np
 import pytorch_lightning as L
-from albumentations import Compose
 from torch.utils.data import DataLoader
 from torch.utils.data.dataloader import default_collate
 
-from careamics.config import InferenceModel
+from careamics.config import InferenceConfig
 from careamics.config.support import SupportedData
 from careamics.config.tile_information import TileInformation
 from careamics.dataset.dataset_utils import (
@@ -38,7 +39,7 @@ def _collate_tiles(batch: List[Tuple[np.ndarray, TileInformation]]) -> Any:
 
     Parameters
     ----------
-    batch : Tuple[Tuple[np.ndarray, TileInformation], ...]
+    batch : List[Tuple[np.ndarray, TileInformation], ...]
         Batch of tiles.
 
     Returns
@@ -62,9 +63,9 @@ def _collate_tiles(batch: List[Tuple[np.ndarray, TileInformation]]) -> Any:
         return default_collate(new_batch)
 
 
-class CAREamicsClay(L.LightningDataModule):
+class CAREamicsPredictData(L.LightningDataModule):
     """
-    LightningDataModule for prediction dataset.
+    CAREamics Lightning prediction data module.
 
     The data module can be used with Path, str or numpy arrays. The data can be either
     a folder containing images or a single file.
@@ -79,7 +80,7 @@ class CAREamicsClay(L.LightningDataModule):
 
     Parameters
     ----------
-    prediction_config : InferenceModel
+    pred_config : InferenceModel
         Pydantic model for CAREamics prediction configuration.
     pred_data : Union[Path, str, np.ndarray]
         Prediction data, can be a path to a folder, a file or a numpy array.
@@ -93,7 +94,7 @@ class CAREamicsClay(L.LightningDataModule):
 
     def __init__(
         self,
-        prediction_config: InferenceModel,
+        pred_config: InferenceConfig,
         pred_data: Union[Path, str, np.ndarray],
         read_source_func: Optional[Callable] = None,
         extension_filter: str = "",
@@ -115,7 +116,7 @@ class CAREamicsClay(L.LightningDataModule):
 
         Parameters
         ----------
-        prediction_config : InferenceModel
+        pred_config : InferenceModel
             Pydantic model for CAREamics prediction configuration.
         pred_data : Union[Path, str, np.ndarray]
             Prediction data, can be a path to a folder, a file or a numpy array.
@@ -142,51 +143,53 @@ class CAREamicsClay(L.LightningDataModule):
         super().__init__()
 
         # check that a read source function is provided for custom types
-        if (
-            prediction_config.data_type == SupportedData.CUSTOM
-            and read_source_func is None
-        ):
+        if pred_config.data_type == SupportedData.CUSTOM and read_source_func is None:
             raise ValueError(
                 f"Data type {SupportedData.CUSTOM} is not allowed without "
-                f"specifying a `read_source_func`."
+                f"specifying a `read_source_func` and an `extension_filer`."
             )
 
-        # and that arrays are passed, if array type specified
-        elif prediction_config.data_type == SupportedData.ARRAY and not isinstance(
-            pred_data, np.ndarray
+        # check correct input type
+        if (
+            isinstance(pred_data, np.ndarray)
+            and pred_config.data_type != SupportedData.ARRAY
         ):
             raise ValueError(
-                f"Expected array input (see configuration.data.data_type), but got "
-                f"{type(pred_data)} instead."
+                f"Received a numpy array as input, but the data type was set to "
+                f"{pred_config.data_type}. Set the data type "
+                f"to {SupportedData.ARRAY} to predict on numpy arrays."
             )
 
         # and that Path or str are passed, if tiff file type specified
-        elif prediction_config.data_type == SupportedData.TIFF and not (
-            isinstance(pred_data, Path) or isinstance(pred_data, str)
+        elif (isinstance(pred_data, Path) or isinstance(pred_config, str)) and (
+            pred_config.data_type != SupportedData.TIFF
+            and pred_config.data_type != SupportedData.CUSTOM
         ):
             raise ValueError(
-                f"Expected Path or str input (see configuration.data.data_type), "
-                f"but got {type(pred_data)} instead."
+                f"Received a path as input, but the data type was neither set to "
+                f"{SupportedData.TIFF} nor {SupportedData.CUSTOM}. Set the data type "
+                f" to {SupportedData.TIFF} or "
+                f"{SupportedData.CUSTOM} to predict on files."
             )
 
         # configuration data
-        self.prediction_config = prediction_config
-        self.data_type = prediction_config.data_type
-        self.batch_size = prediction_config.batch_size
+        self.prediction_config = pred_config
+        self.data_type = pred_config.data_type
+        self.batch_size = pred_config.batch_size
         self.dataloader_params = dataloader_params
 
         self.pred_data = pred_data
-        self.tile_size = prediction_config.tile_size
-        self.tile_overlap = prediction_config.tile_overlap
+        self.tile_size = pred_config.tile_size
+        self.tile_overlap = pred_config.tile_overlap
 
         # read source function
-        if prediction_config.data_type == SupportedData.CUSTOM:
+        if pred_config.data_type == SupportedData.CUSTOM:
             # mypy check
             assert read_source_func is not None
 
             self.read_source_func: Callable = read_source_func
-        elif prediction_config.data_type != SupportedData.ARRAY:
-            self.read_source_func = get_read_func(prediction_config.data_type)
+        elif pred_config.data_type != SupportedData.ARRAY:
+            self.read_source_func = get_read_func(pred_config.data_type)
 
         self.extension_filter = extension_filter
 
@@ -238,9 +241,12 @@ class CAREamicsClay(L.LightningDataModule):
         )  # TODO check workers are used
 
 
-class CAREamicsPredictDataModule(CAREamicsClay):
+class PredictDataWrapper(CAREamicsPredictData):
     """
-    LightningDataModule wrapper of an inference dataset.
+    Wrapper around the CAREamics inference Lightning data module.
+
+    This class is used to explicitely pass the parameters usually contained in a
+    `inference_model` configuration.
 
     Since the lightning datamodule has no access to the model, make sure that the
     parameters passed to the datamodule are consistent with the model's requirements
@@ -251,14 +257,13 @@ class CAREamicsPredictDataModule(CAREamicsClay):
 
     The default transformations applied to the images are defined in
     `careamics.config.inference_model`. To use different transformations, pass a list
-    of transforms or an albumentation `Compose` as `transforms` parameter. See examples
+    of transforms. See examples
     for more details.
 
     The `mean` and `std` parameters are only used if Normalization is defined either
-    in the default transformations or in the `transforms` parameter, but not with
-    a `Compose` object. If you pass a `Normalization` transform in a list as
-    `transforms`, then the mean and std parameters will be overwritten by those passed
-    to this method.
+    in the default transformations or in the `transforms` parameter. If you pass a
+    `Normalization` transform in a list as `transforms`, then the mean and std
+    parameters will be overwritten by those passed to this method.
 
     By default, CAREamics only supports types defined in
     `careamics.config.support.SupportedData`. To read custom data types, you can set
@@ -270,6 +275,12 @@ class CAREamicsPredictDataModule(CAREamicsClay):
     dataloaders, except for `batch_size`, which is set by the `batch_size`
     parameter.
 
+    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
+    avoids artefacts arising from the broken shift invariance induced by the
+    pooling layers of the UNet. If your image has less dimensions, as it may
+    happen in the Z dimension, consider padding your image.
+
     Parameters
     ----------
     pred_data : Union[str, Path, np.ndarray]
@@ -292,7 +303,7 @@ class CAREamicsPredictDataModule(CAREamicsClay):
         Batch size.
     tta_transforms : bool, optional
         Use test time augmentation, by default True.
-    transforms : Optional[Union[List[TRANSFORMS_UNION], Compose]], optional
+    transforms : List, optional
         List of transforms to apply to prediction patches. If None, default
         transforms are applied.
     read_source_func : Optional[Callable], optional
@@ -315,7 +326,7 @@ class CAREamicsPredictDataModule(CAREamicsClay):
         axes: str = "YX",
         batch_size: int = 1,
         tta_transforms: bool = True,
-        transforms: Optional[Union[List, Compose]] = None,
+        transforms: Optional[List] = None,
         read_source_func: Optional[Callable] = None,
         extension_filter: str = "",
         dataloader_params: Optional[dict] = None,
@@ -329,6 +340,12 @@ class CAREamicsPredictDataModule(CAREamicsClay):
             Prediction data.
         data_type : Union[Literal["array", "tiff", "custom"], SupportedData]
             Data type, see `SupportedData` for available options.
+        mean : float
+            Mean value for normalization, only used if Normalization is defined in the
+            transforms.
+        std : float
+            Standard deviation value for normalization, only used if Normalization is
+            defined in the transform.
         tile_size : List[int]
             Tile size, 2D or 3D tile size.
         tile_overlap : List[int]
@@ -339,13 +356,7 @@ class CAREamicsPredictDataModule(CAREamicsClay):
             Batch size.
         tta_transforms : bool, optional
             Use test time augmentation, by default True.
-        mean : Optional[float], optional
-            Mean value for normalization, only used if Normalization is defined, by
-            default None.
-        std : Optional[float], optional
-            Standard deviation value for normalization, only used if Normalization is
-            defined, by default None.
-        transforms : Optional[Union[List[TRANSFORMS_UNION], Compose]], optional
+        transforms : Optional[List], optional
             List of transforms to apply to prediction patches. If None, default
             transforms are applied.
         read_source_func : Optional[Callable], optional
@@ -374,7 +385,7 @@ class CAREamicsPredictDataModule(CAREamicsClay):
             prediction_dict["transforms"] = transforms
 
         # validate configuration
-        self.prediction_config = InferenceModel(**prediction_dict)
+        self.prediction_config = InferenceConfig(**prediction_dict)
 
         # sanity check on the dataloader parameters
         if "batch_size" in dataloader_params:
@@ -382,7 +393,7 @@ class CAREamicsPredictDataModule(CAREamicsClay):
             del dataloader_params["batch_size"]
 
         super().__init__(
-            prediction_config=self.prediction_config,
+            pred_config=self.prediction_config,
             pred_data=pred_data,
             read_source_func=read_source_func,
             extension_filter=extension_filter,

careamics/losses/__init__.py

@@ -1,6 +1,5 @@
 """Losses module."""
 
-
 from .loss_factory import loss_factory
 
 # from .noise_model_factory import noise_model_factory as noise_model_factory

careamics/losses/loss_factory.py

@@ -3,6 +3,7 @@ Loss factory module.
 
 This module contains a factory function for creating loss functions.
 """
+
 from typing import Callable, Union
 
 from ..config.support import SupportedLoss