careamics 0.1.0rc4__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
@@ -73,8 +74,7 @@ class CAREamist:
         source: Union[Path, str],
         work_dir: Optional[str] = None,
         experiment_name: str = "CAREamics",
-    ) -> None:
-        ...
+    ) -> None: ...
 
     @overload
     def __init__(  # numpydoc ignore=GL08
@@ -82,8 +82,7 @@ class CAREamist:
         source: Configuration,
         work_dir: Optional[str] = None,
         experiment_name: str = "CAREamics",
-    ) -> None:
-        ...
+    ) -> None: ...
 
     def __init__(
         self,
@@ -478,8 +477,7 @@ class CAREamist:
         source: CAREamicsPredictData,
         *,
         checkpoint: Optional[Literal["best", "last"]] = None,
-    ) -> Union[list, np.ndarray]:
-        ...
+    ) -> Union[list, np.ndarray]: ...
 
     @overload
     def predict(  # numpydoc ignore=GL08
@@ -491,14 +489,12 @@ 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,
         extension_filter: str = "",
         checkpoint: Optional[Literal["best", "last"]] = None,
-    ) -> Union[list, np.ndarray]:
-        ...
+    ) -> Union[list, np.ndarray]: ...
 
     @overload
     def predict(  # numpydoc ignore=GL08
@@ -510,12 +506,10 @@ 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,
-    ) -> Union[list, np.ndarray]:
-        ...
+    ) -> Union[list, np.ndarray]: ...
 
     def predict(
         self,
@@ -526,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,
@@ -543,11 +536,15 @@ 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.
 
+        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
         ----------
         source : Union[CAREamicsClay, Path, str, np.ndarray]
@@ -562,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
@@ -602,12 +597,11 @@ class CAREamist:
                 )
             # create predict config, reuse training config if parameters missing
             prediction_config = create_inference_configuration(
-                training_configuration=self.cfg,
+                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,
             )
@@ -659,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
@@ -698,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,
@@ -720,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/__init__.py

@@ -1,6 +1,5 @@
 """Configuration module."""
 
-
 __all__ = [
     "AlgorithmConfig",
     "DataConfig",

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
@@ -39,6 +41,7 @@ class UNetModel(ArchitectureModel):
         "None", "Sigmoid", "Softmax", "Tanh", "ReLU", "LeakyReLU"
     ] = Field(default="None", validate_default=True)
     n2v2: bool = Field(default=False, validate_default=True)
+    independent_channels: bool = Field(default=True, validate_default=True)
 
     @field_validator("num_channels_init")
     @classmethod

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,5 @@
-"""Checkpoint saving configuration."""
+"""Callback Pydantic models."""
+
 from __future__ import annotations
 
 from datetime import timedelta
@@ -12,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,
@@ -45,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,12 +58,10 @@ def full_configuration_example() -> Configuration:
                 "name": SupportedTransform.NORMALIZE.value,
             },
             {
-                "name": SupportedTransform.NDFLIP.value,
-                "is_3D": False,
+                "name": SupportedTransform.XY_FLIP.value,
             },
             {
                 "name": SupportedTransform.XY_RANDOM_ROTATE90.value,
-                "is_3D": False,
             },
             {
                 "name": SupportedTransform.N2V_MANIPULATE.value,