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

careamics/callbacks/hyperparameters_callback.py

@@ -1,3 +1,5 @@
+"""Callback saving CAREamics configuration as hyperparameters in the model."""
+
 from pytorch_lightning import LightningModule, Trainer
 from pytorch_lightning.callbacks import Callback
 
@@ -11,13 +13,18 @@ class HyperParametersCallback(Callback):
     This allows saving the configuration as dictionnary in the checkpoints, and
     loading it subsequently in a CAREamist instance.
 
+    Parameters
+    ----------
+    config : Configuration
+        CAREamics configuration to be saved as hyperparameter in the model.
+
     Attributes
     ----------
     config : Configuration
         CAREamics configuration to be saved as hyperparameter in the model.
     """
 
-    def __init__(self, config: Configuration):
+    def __init__(self, config: Configuration) -> None:
         """
         Constructor.
 
@@ -28,14 +35,14 @@ class HyperParametersCallback(Callback):
         """
         self.config = config
 
-    def on_train_start(self, trainer: Trainer, pl_module: LightningModule):
+    def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
         """
         Update the hyperparameters of the model with the configuration on train start.
 
         Parameters
         ----------
         trainer : Trainer
-            PyTorch Lightning trainer.
+            PyTorch Lightning trainer, unused.
         pl_module : LightningModule
             PyTorch Lightning module.
         """

careamics/callbacks/progress_bar_callback.py

@@ -1,3 +1,5 @@
+"""Progressbar callback."""
+
 import sys
 from typing import Dict, Union
 
@@ -10,7 +12,13 @@ class ProgressBarCallback(TQDMProgressBar):
     """Progress bar for training and validation steps."""
 
     def init_train_tqdm(self) -> tqdm:
-        """Override this to customize the tqdm bar for training."""
+        """Override this to customize the tqdm bar for training.
+
+        Returns
+        -------
+        tqdm
+            A tqdm bar.
+        """
         bar = tqdm(
             desc="Training",
             position=(2 * self.process_position),
@@ -23,7 +31,13 @@ class ProgressBarCallback(TQDMProgressBar):
         return bar
 
     def init_validation_tqdm(self) -> tqdm:
-        """Override this to customize the tqdm bar for validation."""
+        """Override this to customize the tqdm bar for validation.
+
+        Returns
+        -------
+        tqdm
+            A tqdm bar.
+        """
         # The main progress bar doesn't exist in `trainer.validate()`
         has_main_bar = self.train_progress_bar is not None
         bar = tqdm(
@@ -37,7 +51,13 @@ class ProgressBarCallback(TQDMProgressBar):
         return bar
 
     def init_test_tqdm(self) -> tqdm:
-        """Override this to customize the tqdm bar for testing."""
+        """Override this to customize the tqdm bar for testing.
+
+        Returns
+        -------
+        tqdm
+            A tqdm bar.
+        """
         bar = tqdm(
             desc="Testing",
             position=(2 * self.process_position),
@@ -52,6 +72,19 @@ class ProgressBarCallback(TQDMProgressBar):
     def get_metrics(
         self, trainer: Trainer, pl_module: LightningModule
     ) -> Dict[str, Union[int, str, float, Dict[str, float]]]:
-        """Override this to customize the metrics displayed in the progress bar."""
+        """Override this to customize the metrics displayed in the progress bar.
+
+        Parameters
+        ----------
+        trainer : Trainer
+            The trainer object.
+        pl_module : LightningModule
+            The LightningModule object, unused.
+
+        Returns
+        -------
+        dict
+            A dictionary with the metrics to display in the progress bar.
+        """
         pbar_metrics = trainer.progress_bar_metrics
         return {**pbar_metrics}

careamics/careamist.py

@@ -18,13 +18,14 @@ from careamics.config import (
     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.transforms import Denormalize
 from careamics.utils import check_path_exists, get_logger
 
 from .callbacks import HyperParametersCallback
@@ -488,7 +489,6 @@ class CAREamist:
         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,
         read_source_func: Optional[Callable] = None,
@@ -506,7 +506,6 @@ class CAREamist:
         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,
         checkpoint: Optional[Literal["best", "last"]] = None,
@@ -521,7 +520,6 @@ class CAREamist:
         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,
         read_source_func: Optional[Callable] = None,
@@ -538,8 +536,6 @@ class CAREamist:
         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.
 
@@ -563,8 +559,6 @@ class CAREamist:
             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
@@ -608,7 +602,6 @@ class CAREamist:
                 tile_overlap=tile_overlap,
                 data_type=data_type,
                 axes=axes,
-                transforms=transforms,
                 tta_transforms=tta_transforms,
                 batch_size=batch_size,
             )
@@ -660,38 +653,41 @@ class CAREamist:
                     f"np.ndarray (got {type(source)})."
                 )
 
-    def export_to_bmz(
+    def _create_data_for_bmz(
         self,
-        path: Union[Path, str],
-        name: str,
-        authors: List[dict],
         input_array: Optional[np.ndarray] = None,
-        general_description: str = "",
-        channel_names: Optional[List[str]] = None,
-        data_description: Optional[str] = None,
-    ) -> None:
-        """Export the model to the BioImage Model Zoo format.
+    ) -> np.ndarray:
+        """Create data for BMZ export.
 
-        Input array must be of shape SC(Z)YX, with S and C singleton dimensions.
+        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
         ----------
-        path : Union[Path, str]
-            Path to save the model.
-        name : str
-            Name of the model.
-        authors : List[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, by default None.
-        data_description : Optional[str], optional
-            Description of the data, by default None.
+            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.
         """
         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
@@ -699,19 +695,23 @@ class CAREamist:
 
                 # 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()
-            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?"
-                    )
 
+                # 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,
@@ -721,11 +721,47 @@ class CAREamist:
                     np.newaxis, np.newaxis, ...
                 ]  # add S & C dimensions
         else:
-            input_patch = input_array
+            # potentially correct shape
+            input_patch = reshape_array(input_array, self.cfg.data_config.axes)
 
-        # if there is a batch dimension
+        # if this a batch
         if input_patch.shape[0] > 1:
-            input_patch = input_patch[0:1, ...]  # keep singleton dim
+            input_patch = input_patch[[0], ...]  # keep singleton dim
+
+        return input_patch
+
+    def export_to_bmz(
+        self,
+        path: Union[Path, str],
+        name: str,
+        authors: List[dict],
+        input_array: Optional[np.ndarray] = None,
+        general_description: str = "",
+        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.
+
+        Parameters
+        ----------
+        path : Union[Path, str]
+            Path to save the model.
+        name : str
+            Name of the model.
+        authors : List[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, by default None.
+        data_description : Optional[str], optional
+            Description of the data, by default None.
+        """
+        input_patch = self._create_data_for_bmz(input_array)
 
         # axes need to be reformated for the export because reshaping was done in the
         # datamodule

careamics/config/algorithm_model.py

@@ -1,3 +1,5 @@
+"""Algorithm configuration."""
+
 from __future__ import annotations
 
 from pprint import pformat
@@ -17,9 +19,9 @@ class AlgorithmConfig(BaseModel):
     training algorithm: which algorithm, loss function, model architecture, optimizer,
     and learning rate scheduler to use.
 
-    Currently, we only support N2V and custom algorithms. The `n2v` algorithm is only
-    compatible with `n2v` loss and `UNet` architecture. The `custom` algorithm allows
-    you to register your own architecture and select it using its name as
+    Currently, we only support N2V, CARE, N2N and custom models. The `n2v` algorithm is
+    only compatible with `n2v` loss and `UNet` architecture. The `custom` algorithm
+    allows you to register your own architecture and select it using its name as
     `name` in the custom pydantic model.
 
     Attributes

careamics/config/architectures/architecture_model.py

@@ -1,3 +1,5 @@
+"""Base model for the various CAREamics architectures."""
+
 from typing import Any, Dict
 
 from pydantic import BaseModel
@@ -16,6 +18,11 @@ class ArchitectureModel(BaseModel):
         """
         Dump the model as a dictionary, ignoring the architecture keyword.
 
+        Parameters
+        ----------
+        **kwargs : Any
+            Additional keyword arguments from Pydantic BaseModel model_dump method.
+
         Returns
         -------
         dict[str, Any]

careamics/config/architectures/custom_model.py

@@ -1,3 +1,5 @@
+"""Custom architecture Pydantic model."""
+
 from __future__ import annotations
 
 from pprint import pformat
@@ -84,6 +86,11 @@ class CustomModel(ArchitectureModel):
         value : str
             Name of the custom model as registered using the `@register_model`
             decorator.
+
+        Returns
+        -------
+        str
+            The custom model name.
         """
         # delegate error to get_custom_model
         model = get_custom_model(value)
@@ -134,7 +141,7 @@ class CustomModel(ArchitectureModel):
 
         Parameters
         ----------
-        kwargs : Any
+        **kwargs : Any
             Additional keyword arguments from Pydantic BaseModel model_dump method.
 
         Returns

careamics/config/architectures/register_model.py

@@ -1,3 +1,5 @@
+"""Custom model registration utilities."""
+
 from typing import Callable
 
 from torch.nn import Module
@@ -53,7 +55,7 @@ def register_model(name: str) -> Callable:
         Parameters
         ----------
         model : Module
-            Module class to register
+            Module class to register.
 
         Returns
         -------

careamics/config/architectures/unet_model.py

@@ -1,3 +1,5 @@
+"""UNet Pydantic model."""
+
 from __future__ import annotations
 
 from typing import Literal

careamics/config/architectures/vae_model.py

@@ -1,3 +1,5 @@
+"""VAE Pydantic model."""
+
 from typing import Literal
 
 from pydantic import (

careamics/config/callback_model.py

@@ -1,4 +1,4 @@
-"""Checkpoint saving configuration."""
+"""Callback Pydantic models."""
 
 from __future__ import annotations
 
@@ -13,13 +13,7 @@ from pydantic import (
 
 
 class CheckpointModel(BaseModel):
-    """_summary_.
-
-    Parameters
-    ----------
-    BaseModel : _type_
-        _description_
-    """
+    """Checkpoint saving callback Pydantic model."""
 
     model_config = ConfigDict(
         validate_assignment=True,
@@ -46,13 +40,7 @@ class CheckpointModel(BaseModel):
 
 
 class EarlyStoppingModel(BaseModel):
-    """_summary_.
-
-    Parameters
-    ----------
-    BaseModel : _type_
-        _description_
-    """
+    """Early stopping callback Pydantic model."""
 
     model_config = ConfigDict(
         validate_assignment=True,

careamics/config/configuration_example.py

@@ -1,3 +1,5 @@
+"""Example of configurations."""
+
 from .algorithm_model import AlgorithmConfig
 from .architectures import UNetModel
 from .configuration_model import Configuration
@@ -19,7 +21,7 @@ from .training_model import TrainingConfig
 
 
 def full_configuration_example() -> Configuration:
-    """Returns a dictionnary representing a full configuration example.
+    """Return a dictionnary representing a full configuration example.
 
     Returns
     -------
@@ -56,7 +58,7 @@ def full_configuration_example() -> Configuration:
                 "name": SupportedTransform.NORMALIZE.value,
             },
             {
-                "name": SupportedTransform.NDFLIP.value,
+                "name": SupportedTransform.XY_FLIP.value,
             },
             {
                 "name": SupportedTransform.XY_RANDOM_ROTATE90.value,

careamics/config/configuration_factory.py

@@ -1,6 +1,6 @@
 """Convenience functions to create configurations for training and inference."""
 
-from typing import Any, Dict, List, Literal, Optional, Tuple, Union
+from typing import Any, Dict, List, Literal, Optional, Tuple
 
 from .algorithm_model import AlgorithmConfig
 from .architectures import UNetModel
@@ -111,7 +111,7 @@ def _create_supervised_configuration(
                 "name": SupportedTransform.NORMALIZE.value,
             },
             {
-                "name": SupportedTransform.NDFLIP.value,
+                "name": SupportedTransform.XY_FLIP.value,
             },
             {
                 "name": SupportedTransform.XY_RANDOM_ROTATE90.value,
@@ -526,7 +526,7 @@ def create_n2v_configuration(
                 "name": SupportedTransform.NORMALIZE.value,
             },
             {
-                "name": SupportedTransform.NDFLIP.value,
+                "name": SupportedTransform.XY_FLIP.value,
             },
             {
                 "name": SupportedTransform.XY_RANDOM_ROTATE90.value,
@@ -587,7 +587,6 @@ def create_inference_configuration(
     tile_overlap: Optional[Tuple[int, ...]] = None,
     data_type: Optional[Literal["array", "tiff", "custom"]] = None,
     axes: Optional[str] = None,
-    transforms: Optional[Union[List[Dict[str, Any]]]] = None,
     tta_transforms: bool = True,
     batch_size: Optional[int] = 1,
 ) -> InferenceConfig:
@@ -595,7 +594,7 @@ def create_inference_configuration(
     Create a configuration for inference with N2V.
 
     If not provided, `data_type` and `axes` are taken from the training
-    configuration. If `transforms` are not provided, only normalization is applied.
+    configuration.
 
     Parameters
     ----------
@@ -609,8 +608,6 @@ def create_inference_configuration(
         Type of the data, by default "tiff".
     axes : str, optional
         Axes of the data, by default "YX".
-    transforms : List[Dict[str, Any]], optional
-        Transformations to apply to the data, by default None.
     tta_transforms : bool, optional
         Whether to apply test-time augmentations, by default True.
     batch_size : int, optional
@@ -624,14 +621,6 @@ def create_inference_configuration(
     if configuration.data_config.mean is None or configuration.data_config.std is None:
         raise ValueError("Mean and std must be provided in the configuration.")
 
-    # minimum transform
-    if transforms is None:
-        transforms = [
-            {
-                "name": SupportedTransform.NORMALIZE.value,
-            },
-        ]
-
     # tile size for UNets
     if tile_size is not None:
         model = configuration.algorithm_config.model
@@ -661,7 +650,6 @@ def create_inference_configuration(
         axes=axes or configuration.data_config.axes,
         mean=configuration.data_config.mean,
         std=configuration.data_config.std,
-        transforms=transforms,
         tta_transforms=tta_transforms,
         batch_size=batch_size,
     )

careamics/config/data_model.py

@@ -17,14 +17,14 @@ from typing_extensions import Annotated, Self
 
 from .support import SupportedTransform
 from .transformations.n2v_manipulate_model import N2VManipulateModel
-from .transformations.nd_flip_model import NDFlipModel
 from .transformations.normalize_model import NormalizeModel
+from .transformations.xy_flip_model import XYFlipModel
 from .transformations.xy_random_rotate90_model import XYRandomRotate90Model
 from .validators import check_axes_validity, patch_size_ge_than_8_power_of_2
 
 TRANSFORMS_UNION = Annotated[
     Union[
-        NDFlipModel,
+        XYFlipModel,
         XYRandomRotate90Model,
         NormalizeModel,
         N2VManipulateModel,
@@ -41,6 +41,8 @@ class DataConfig(BaseModel):
     and then the mean (if they were both `None` before) will raise a validation error.
     Prefer instead `set_mean_and_std` to set both at once.
 
+    All supported transforms are defined in the SupportedTransform enum.
+
     Examples
     --------
     Minimum example:
@@ -56,7 +58,7 @@ class DataConfig(BaseModel):
     >>> data.set_mean_and_std(mean=214.3, std=84.5)
 
     One can pass also a list of transformations, by keyword, using the
-    SupportedTransform or the name of an Albumentation transform:
+    SupportedTransform value:
     >>> from careamics.config.support import SupportedTransform
     >>> data = DataConfig(
     ...     data_type="tiff",
@@ -70,7 +72,7 @@ class DataConfig(BaseModel):
     ...             "std": 47.2,
     ...         },
     ...         {
-    ...             "name": "NDFlip",
+    ...             "name": "XYFlip",
     ...         }
     ...     ]
     ... )
@@ -97,7 +99,7 @@ class DataConfig(BaseModel):
                 "name": SupportedTransform.NORMALIZE.value,
             },
             {
-                "name": SupportedTransform.NDFLIP.value,
+                "name": SupportedTransform.XY_FLIP.value,
             },
             {
                 "name": SupportedTransform.XY_RANDOM_ROTATE90.value,
@@ -202,7 +204,7 @@ class DataConfig(BaseModel):
 
         if SupportedTransform.N2V_MANIPULATE in transform_list:
             # multiple N2V_MANIPULATE
-            if transform_list.count(SupportedTransform.N2V_MANIPULATE) > 1:
+            if transform_list.count(SupportedTransform.N2V_MANIPULATE.value) > 1:
                 raise ValueError(
                     f"Multiple instances of "
                     f"{SupportedTransform.N2V_MANIPULATE} transforms "
@@ -211,7 +213,7 @@ class DataConfig(BaseModel):
 
             # N2V_MANIPULATE not the last transform
             elif transform_list[-1] != SupportedTransform.N2V_MANIPULATE:
-                index = transform_list.index(SupportedTransform.N2V_MANIPULATE)
+                index = transform_list.index(SupportedTransform.N2V_MANIPULATE.value)
                 transform = transforms.pop(index)
                 transforms.append(transform)
 
@@ -250,7 +252,7 @@ class DataConfig(BaseModel):
         Self
             Data model with mean and std added to the Normalize transform.
         """
-        if self.mean is not None or self.std is not None:
+        if self.mean is not None and self.std is not None:
             # search in the transforms for Normalize and update parameters
             for transform in self.transforms:
                 if transform.name == SupportedTransform.NORMALIZE.value:
@@ -355,12 +357,6 @@ class DataConfig(BaseModel):
         """
         self._update(mean=mean, std=std)
 
-        # search in the transforms for Normalize and update parameters
-        for transform in self.transforms:
-            if transform.name == SupportedTransform.NORMALIZE.value:
-                transform.mean = mean
-                transform.std = std
-
     def set_3D(self, axes: str, patch_size: List[int]) -> None:
         """
         Set 3D parameters.