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

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
@@ -132,21 +134,6 @@ class AlgorithmConfig(BaseModel):
                     "sure that `in_channels` and `num_classes` are the same."
                 )
 
-        # N2N
-        if self.algorithm == "n2n":
-            # n2n is only compatible with the UNet model
-            if not isinstance(self.model, UNetModel):
-                raise ValueError(
-                    f"Model for algorithm {self.algorithm} must be a `UNetModel`."
-                )
-
-            # n2n requires the number of input and output channels to be the same
-            if self.model.in_channels != self.model.num_classes:
-                raise ValueError(
-                    "N2N requires the same number of input and output channels. Make "
-                    "sure that `in_channels` and `num_classes` are the same."
-                )
-
         if self.algorithm == "care" or self.algorithm == "n2n":
             if self.loss == "n2v":
                 raise ValueError("Supervised algorithms do not support loss `n2v`.")

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,7 +1,9 @@
+"""Custom architecture Pydantic model."""
+
 from __future__ import annotations
 
 from pprint import pformat
-from typing import Any, Dict, Literal
+from typing import Any, Literal
 
 from pydantic import ConfigDict, field_validator, model_validator
 from torch.nn import Module
@@ -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)
@@ -129,17 +136,17 @@ class CustomModel(ArchitectureModel):
         """
         return pformat(self.model_dump())
 
-    def model_dump(self, **kwargs: Any) -> Dict[str, Any]:
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
         """Dump the model configuration.
 
         Parameters
         ----------
-        kwargs : Any
+        **kwargs : Any
             Additional keyword arguments from Pydantic BaseModel model_dump method.
 
         Returns
         -------
-        Dict[str, Any]
+        dict[str, Any]
             Model configuration.
         """
         model_dict = super().model_dump()

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
     -------
@@ -53,10 +55,7 @@ def full_configuration_example() -> Configuration:
         axes="YX",
         transforms=[
             {
-                "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
@@ -108,21 +108,14 @@ def _create_supervised_configuration(
     if use_augmentations:
         transforms: List[Dict[str, Any]] = [
             {
-                "name": SupportedTransform.NORMALIZE.value,
-            },
-            {
-                "name": SupportedTransform.NDFLIP.value,
+                "name": SupportedTransform.XY_FLIP.value,
             },
             {
                 "name": SupportedTransform.XY_RANDOM_ROTATE90.value,
             },
         ]
     else:
-        transforms = [
-            {
-                "name": SupportedTransform.NORMALIZE.value,
-            },
-        ]
+        transforms = []
 
     # data model
     data = DataConfig(
@@ -250,7 +243,8 @@ def create_n2n_configuration(
     use_augmentations: bool = True,
     independent_channels: bool = False,
     loss: Literal["mae", "mse"] = "mae",
-    n_channels: int = 1,
+    n_channels_in: int = 1,
+    n_channels_out: int = -1,
     logger: Literal["wandb", "tensorboard", "none"] = "none",
     model_kwargs: Optional[dict] = None,
 ) -> Configuration:
@@ -260,10 +254,13 @@ def create_n2n_configuration(
     If "Z" is present in `axes`, then `path_size` must be a list of length 3, otherwise
     2.
 
-    If "C" is present in `axes`, then you need to set `n_channels` to the number of
+    If "C" is present in `axes`, then you need to set `n_channels_in` to the number of
     channels. Likewise, if you set the number of channels, then "C" must be present in
     `axes`.
 
+    To set the number of output channels, use the `n_channels_out` parameter. If it is
+    not specified, it will be assumed to be equal to `n_channels_in`.
+
     By default, all channels are trained together. To train all channels independently,
     set `independent_channels` to True.
 
@@ -290,8 +287,10 @@ def create_n2n_configuration(
         Whether to train all channels independently, by default False.
     loss : Literal["mae", "mse"], optional
         Loss function to use, by default "mae".
-    n_channels : int, optional
-        Number of channels (in and out), by default 1.
+    n_channels_in : int, optional
+        Number of channels in, by default 1.
+    n_channels_out : int, optional
+        Number of channels out, by default -1.
     logger : Literal["wandb", "tensorboard", "none"], optional
         Logger to use, by default "none".
     model_kwargs : dict, optional
@@ -302,6 +301,9 @@ def create_n2n_configuration(
     Configuration
         Configuration for training Noise2Noise.
     """
+    if n_channels_out == -1:
+        n_channels_out = n_channels_in
+
     return _create_supervised_configuration(
         algorithm="n2n",
         experiment_name=experiment_name,
@@ -313,8 +315,8 @@ def create_n2n_configuration(
         use_augmentations=use_augmentations,
         independent_channels=independent_channels,
         loss=loss,
-        n_channels_in=n_channels,
-        n_channels_out=n_channels,
+        n_channels_in=n_channels_in,
+        n_channels_out=n_channels_out,
         logger=logger,
         model_kwargs=model_kwargs,
     )
@@ -523,21 +525,14 @@ def create_n2v_configuration(
     if use_augmentations:
         transforms: List[Dict[str, Any]] = [
             {
-                "name": SupportedTransform.NORMALIZE.value,
-            },
-            {
-                "name": SupportedTransform.NDFLIP.value,
+                "name": SupportedTransform.XY_FLIP.value,
             },
             {
                 "name": SupportedTransform.XY_RANDOM_ROTATE90.value,
             },
         ]
     else:
-        transforms = [
-            {
-                "name": SupportedTransform.NORMALIZE.value,
-            },
-        ]
+        transforms = []
 
     # n2v2 and structn2v
     nv2_transform = {
@@ -587,7 +582,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 +589,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 +603,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
@@ -621,17 +613,12 @@ def create_inference_configuration(
     InferenceConfiguration
         Configuration used to configure CAREamicsPredictData.
     """
-    if configuration.data_config.mean is None or configuration.data_config.std is None:
+    if (
+        configuration.data_config.image_means is None
+        or configuration.data_config.image_stds 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
@@ -659,9 +646,8 @@ def create_inference_configuration(
         tile_size=tile_size,
         tile_overlap=tile_overlap,
         axes=axes or configuration.data_config.axes,
-        mean=configuration.data_config.mean,
-        std=configuration.data_config.std,
-        transforms=transforms,
+        image_means=configuration.data_config.image_means,
+        image_stds=configuration.data_config.image_stds,
         tta_transforms=tta_transforms,
         batch_size=batch_size,
     )

careamics/config/configuration_model.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 import re
 from pathlib import Path
 from pprint import pformat
-from typing import Dict, List, Literal, Union
+from typing import Literal, Union
 
 import yaml
 from bioimageio.spec.generic.v0_3 import CiteEntry
@@ -269,7 +269,7 @@ class Configuration(BaseModel):
         """
         return pformat(self.model_dump())
 
-    def set_3D(self, is_3D: bool, axes: str, patch_size: List[int]) -> None:
+    def set_3D(self, is_3D: bool, axes: str, patch_size: list[int]) -> None:
         """
         Set 3D flag and axes.
 
@@ -279,7 +279,7 @@ class Configuration(BaseModel):
             Whether the algorithm is 3D or not.
         axes : str
             Axes of the data.
-        patch_size : List[int]
+        patch_size : list[int]
             Patch size.
         """
         # set the flag and axes (this will not trigger validation at the config level)
@@ -389,7 +389,7 @@ class Configuration(BaseModel):
 
         return ""
 
-    def get_algorithm_citations(self) -> List[CiteEntry]:
+    def get_algorithm_citations(self) -> list[CiteEntry]:
         """
         Return a list of citation entries of the current algorithm.
 
@@ -455,13 +455,13 @@ class Configuration(BaseModel):
 
         return ""
 
-    def get_algorithm_keywords(self) -> List[str]:
+    def get_algorithm_keywords(self) -> list[str]:
         """
         Get algorithm keywords.
 
         Returns
         -------
-        List[str]
+        list[str]
             List of keywords.
         """
         if self.algorithm_config.algorithm == SupportedAlgorithm.N2V:
@@ -491,8 +491,8 @@ class Configuration(BaseModel):
         self,
         exclude_defaults: bool = False,
         exclude_none: bool = True,
-        **kwargs: Dict,
-    ) -> Dict:
+        **kwargs: dict,
+    ) -> dict:
         """
         Override model_dump method in order to set default values.
 
@@ -503,7 +503,7 @@ class Configuration(BaseModel):
             True.
         exclude_none : bool, optional
             Whether to exclude fields with None values or not, by default True.
-        **kwargs : Dict
+        **kwargs : dict
             Keyword arguments.
 
         Returns
@@ -524,7 +524,7 @@ def load_configuration(path: Union[str, Path]) -> Configuration:
 
     Parameters
     ----------
-    path : Union[str, Path]
+    path : str or Path
         Path to the configuration.
 
     Returns
@@ -556,7 +556,7 @@ def save_configuration(config: Configuration, path: Union[str, Path]) -> Path:
     ----------
     config : Configuration
         Configuration to save.
-    path : Union[str, Path]
+    path : str or Path
         Path to a existing folder in which to save the configuration or to an existing
         configuration file.