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

careamics/config/data_model.py

@@ -3,8 +3,9 @@
 from __future__ import annotations
 
 from pprint import pformat
-from typing import Any, List, Literal, Optional, Union
+from typing import Any, Literal, Optional, Union
 
+from numpy.typing import NDArray
 from pydantic import (
     BaseModel,
     ConfigDict,
@@ -17,16 +18,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,
     ],
     Discriminator("name"),  # used to tell the different transform models apart
@@ -39,7 +38,11 @@ class DataConfig(BaseModel):
 
     If std is specified, mean must be specified as well. Note that setting the std first
     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.
+    Prefer instead `set_mean_and_std` to set both at once. Means and stds are expected
+    to be lists of floats, one for each channel. For supervised tasks, the mean and std
+    of the target could be different from the input data.
+
+    All supported transforms are defined in the SupportedTransform enum.
 
     Examples
     --------
@@ -53,10 +56,10 @@ class DataConfig(BaseModel):
     ... )
 
     To change the mean and std of the data:
-    >>> data.set_mean_and_std(mean=214.3, std=84.5)
+    >>> data.set_mean_and_std(image_means=[214.3], image_stds=[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",
@@ -65,12 +68,7 @@ class DataConfig(BaseModel):
     ...     axes="YX",
     ...     transforms=[
     ...         {
-    ...             "name": SupportedTransform.NORMALIZE.value,
-    ...             "mean": 167.6,
-    ...             "std": 47.2,
-    ...         },
-    ...         {
-    ...             "name": "NDFlip",
+    ...             "name": "XYFlip",
     ...         }
     ...     ]
     ... )
@@ -83,21 +81,26 @@ class DataConfig(BaseModel):
 
     # Dataset configuration
     data_type: Literal["array", "tiff", "custom"]  # As defined in SupportedData
-    patch_size: Union[List[int]] = Field(..., min_length=2, max_length=3)
+    patch_size: Union[list[int]] = Field(..., min_length=2, max_length=3)
     batch_size: int = Field(default=1, ge=1, validate_default=True)
     axes: str
 
     # Optional fields
-    mean: Optional[float] = None
-    std: Optional[float] = None
+    image_means: Optional[list[float]] = Field(
+        default=None, min_length=0, max_length=32
+    )
+    image_stds: Optional[list[float]] = Field(default=None, min_length=0, max_length=32)
+    target_means: Optional[list[float]] = Field(
+        default=None, min_length=0, max_length=32
+    )
+    target_stds: Optional[list[float]] = Field(
+        default=None, min_length=0, max_length=32
+    )
 
-    transforms: List[TRANSFORMS_UNION] = Field(
+    transforms: list[TRANSFORMS_UNION] = Field(
         default=[
             {
-                "name": SupportedTransform.NORMALIZE.value,
-            },
-            {
-                "name": SupportedTransform.NDFLIP.value,
+                "name": SupportedTransform.XY_FLIP.value,
             },
             {
                 "name": SupportedTransform.XY_RANDOM_ROTATE90.value,
@@ -114,8 +117,8 @@ class DataConfig(BaseModel):
     @field_validator("patch_size")
     @classmethod
     def all_elements_power_of_2_minimum_8(
-        cls, patch_list: Union[List[int]]
-    ) -> Union[List[int]]:
+        cls, patch_list: Union[list[int]]
+    ) -> Union[list[int]]:
         """
         Validate patch size.
 
@@ -123,12 +126,12 @@ class DataConfig(BaseModel):
 
         Parameters
         ----------
-        patch_list : Union[List[int]]
+        patch_list : list of int
             Patch size.
 
         Returns
         -------
-        Union[List[int]]
+        list of int
             Validated patch size.
 
         Raises
@@ -178,19 +181,19 @@ class DataConfig(BaseModel):
     @field_validator("transforms")
     @classmethod
     def validate_prediction_transforms(
-        cls, transforms: List[TRANSFORMS_UNION]
-    ) -> List[TRANSFORMS_UNION]:
+        cls, transforms: list[TRANSFORMS_UNION]
+    ) -> list[TRANSFORMS_UNION]:
         """
         Validate N2VManipulate transform position in the transform list.
 
         Parameters
         ----------
-        transforms : List[Transformations_Union]
+        transforms : list[Transformations_Union]
             Transforms.
 
         Returns
         -------
-        List[TRANSFORMS_UNION]
+        list of transforms
             Validated transforms.
 
         Raises
@@ -202,7 +205,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 +214,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)
 
@@ -233,29 +236,33 @@ class DataConfig(BaseModel):
             If std is not None and mean is None.
         """
         # check that mean and std are either both None, or both specified
-        if (self.mean is None) != (self.std is None):
+        if (self.image_means and not self.image_stds) or (
+            self.image_stds and not self.image_means
+        ):
             raise ValueError(
                 "Mean and std must be either both None, or both specified."
             )
 
-        return self
+        elif (self.image_means is not None and self.image_stds is not None) and (
+            len(self.image_means) != len(self.image_stds)
+        ):
+            raise ValueError(
+                "Mean and std must be specified for each " "input channel."
+            )
 
-    @model_validator(mode="after")
-    def add_std_and_mean_to_normalize(self: Self) -> Self:
-        """
-        Add mean and std to the Normalize transform if it is present.
+        if (self.target_means and not self.target_stds) or (
+            self.target_stds and not self.target_means
+        ):
+            raise ValueError(
+                "Mean and std must be either both None, or both specified "
+            )
 
-        Returns
-        -------
-        Self
-            Data model with mean and std added to the Normalize transform.
-        """
-        if self.mean is not None or 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:
-                    transform.mean = self.mean
-                    transform.std = self.std
+        elif self.target_means is not None and self.target_stds is not None:
+            if len(self.target_means) != len(self.target_stds):
+                raise ValueError(
+                    "Mean and std must be either both None, or both specified for each "
+                    "target channel."
+                )
 
         return self
 
@@ -339,7 +346,13 @@ class DataConfig(BaseModel):
         if self.has_n2v_manipulate():
             self.transforms.pop(-1)
 
-    def set_mean_and_std(self, mean: float, std: float) -> None:
+    def set_mean_and_std(
+        self,
+        image_means: Union[NDArray, tuple, list, None],
+        image_stds: Union[NDArray, tuple, list, None],
+        target_means: Optional[Union[NDArray, tuple, list, None]] = None,
+        target_stds: Optional[Union[NDArray, tuple, list, None]] = None,
+    ) -> None:
         """
         Set mean and standard deviation of the data.
 
@@ -348,20 +361,33 @@ class DataConfig(BaseModel):
 
         Parameters
         ----------
-        mean : float
-            Mean of the data.
-        std : float
-            Standard deviation of the data.
+        image_means : NDArray or tuple or list
+            Mean values for normalization.
+        image_stds : NDArray or tuple or list
+            Standard deviation values for normalization.
+        target_means : NDArray or tuple or list, optional
+            Target mean values for normalization, by default ().
+        target_stds : NDArray or tuple or list, optional
+            Target standard deviation values for normalization, by default ().
         """
-        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
+        # make sure we pass a list
+        if image_means is not None:
+            image_means = list(image_means)
+        if image_stds is not None:
+            image_stds = list(image_stds)
+        if target_means is not None:
+            target_means = list(target_means)
+        if target_stds is not None:
+            target_stds = list(target_stds)
+
+        self._update(
+            image_means=image_means,
+            image_stds=image_stds,
+            target_means=target_means,
+            target_stds=target_stds,
+        )
 
-    def set_3D(self, axes: str, patch_size: List[int]) -> None:
+    def set_3D(self, axes: str, patch_size: list[int]) -> None:
         """
         Set 3D parameters.
 
@@ -369,7 +395,7 @@ class DataConfig(BaseModel):
         ----------
         axes : str
             Axes.
-        patch_size : List[int]
+        patch_size : list of int
             Patch size.
         """
         self._update(axes=axes, patch_size=patch_size)

careamics/config/inference_model.py

@@ -2,17 +2,13 @@
 
 from __future__ import annotations
 
-from typing import Any, List, Literal, Optional, Union
+from typing import Any, Literal, Optional, Union
 
 from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 from typing_extensions import Self
 
-from .support import SupportedTransform
-from .transformations.normalize_model import NormalizeModel
 from .validators import check_axes_validity, patch_size_ge_than_8_power_of_2
 
-TRANSFORMS_UNION = Union[NormalizeModel]
-
 
 class InferenceConfig(BaseModel):
     """Configuration class for the prediction model."""
@@ -21,26 +17,17 @@ class InferenceConfig(BaseModel):
 
     # Mandatory fields
     data_type: Literal["array", "tiff", "custom"]  # As defined in SupportedData
-    tile_size: Optional[Union[List[int]]] = Field(
+    tile_size: Optional[Union[list[int]]] = Field(
         default=None, min_length=2, max_length=3
     )
-    tile_overlap: Optional[Union[List[int]]] = Field(
+    tile_overlap: Optional[Union[list[int]]] = Field(
         default=None, min_length=2, max_length=3
     )
 
     axes: str
 
-    mean: float
-    std: float = Field(..., ge=0.0)
-
-    transforms: List[TRANSFORMS_UNION] = Field(
-        default=[
-            {
-                "name": SupportedTransform.NORMALIZE.value,
-            },
-        ],
-        validate_default=True,
-    )
+    image_means: list = Field(..., min_length=0, max_length=32)
+    image_stds: list = Field(..., min_length=0, max_length=32)
 
     # only default TTAs are supported for now
     tta_transforms: bool = Field(default=True)
@@ -51,8 +38,8 @@ class InferenceConfig(BaseModel):
     @field_validator("tile_overlap")
     @classmethod
     def all_elements_non_zero_even(
-        cls, tile_overlap: Optional[Union[List[int]]]
-    ) -> Optional[Union[List[int]]]:
+        cls, tile_overlap: Optional[list[int]]
+    ) -> Optional[list[int]]:
         """
         Validate tile overlap.
 
@@ -60,12 +47,12 @@ class InferenceConfig(BaseModel):
 
         Parameters
         ----------
-        tile_overlap : Optional[Union[List[int]]]
+        tile_overlap : list[int] or None
             Patch size.
 
         Returns
         -------
-        Optional[Union[List[int]]]
+        list[int] or None
             Validated tile overlap.
 
         Raises
@@ -90,19 +77,19 @@ class InferenceConfig(BaseModel):
     @field_validator("tile_size")
     @classmethod
     def tile_min_8_power_of_2(
-        cls, tile_list: Optional[Union[List[int]]]
-    ) -> Optional[Union[List[int]]]:
+        cls, tile_list: Optional[list[int]]
+    ) -> Optional[list[int]]:
         """
         Validate that each entry is greater or equal than 8 and a power of 2.
 
         Parameters
         ----------
-        tile_list : List[int]
+        tile_list : list of int
             Patch size.
 
         Returns
         -------
-        List[int]
+        list of int
             Validated patch size.
 
         Raises
@@ -149,39 +136,6 @@ class InferenceConfig(BaseModel):
 
         return axes
 
-    @field_validator("transforms")
-    @classmethod
-    def validate_transforms(
-        cls, transforms: List[TRANSFORMS_UNION]
-    ) -> List[TRANSFORMS_UNION]:
-        """
-        Validate that transforms do not have N2V pixel manipulate transforms.
-
-        Parameters
-        ----------
-        transforms : List[TRANSFORMS_UNION]
-            Transforms.
-
-        Returns
-        -------
-        List[TRANSFORMS_UNION]
-            Validated transforms.
-
-        Raises
-        ------
-        ValueError
-            If transforms contain N2V pixel manipulate transforms.
-        """
-        if transforms is not None:
-            for transform in transforms:
-                if transform.name == SupportedTransform.N2V_MANIPULATE.value:
-                    raise ValueError(
-                        "N2V_Manipulate transform is not allowed in "
-                        "prediction transforms."
-                    )
-
-        return transforms
-
     @model_validator(mode="after")
     def validate_dimensions(self: Self) -> Self:
         """
@@ -228,29 +182,22 @@ class InferenceConfig(BaseModel):
             If std is not None and mean is None.
         """
         # check that mean and std are either both None, or both specified
-        if (self.mean is None) != (self.std is None):
+        if not self.image_means and not self.image_stds:
+            raise ValueError("Mean and std must be specified during inference.")
+
+        if (self.image_means and not self.image_stds) or (
+            self.image_stds and not self.image_means
+        ):
             raise ValueError(
                 "Mean and std must be either both None, or both specified."
             )
 
-        return self
-
-    @model_validator(mode="after")
-    def add_std_and_mean_to_normalize(self: Self) -> Self:
-        """
-        Add mean and std to the Normalize transform if it is present.
-
-        Returns
-        -------
-        Self
-            Inference model with mean and std added to the Normalize transform.
-        """
-        if self.mean is not None or 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:
-                    transform.mean = self.mean
-                    transform.std = self.std
+        elif (self.image_means is not None and self.image_stds is not None) and (
+            len(self.image_means) != len(self.image_stds)
+        ):
+            raise ValueError(
+                "Mean and std must be specified for each " "input channel."
+            )
 
         return self
 
@@ -266,7 +213,7 @@ class InferenceConfig(BaseModel):
         self.__dict__.update(kwargs)
         self.__class__.model_validate(self.__dict__)
 
-    def set_3D(self, axes: str, tile_size: List[int], tile_overlap: List[int]) -> None:
+    def set_3D(self, axes: str, tile_size: list[int], tile_overlap: list[int]) -> None:
         """
         Set 3D parameters.
 
@@ -274,9 +221,9 @@ class InferenceConfig(BaseModel):
         ----------
         axes : str
             Axes.
-        tile_size : List[int]
+        tile_size : list of int
             Tile size.
-        tile_overlap : List[int]
+        tile_overlap : list of int
             Tile overlap.
         """
         self._update(axes=axes, tile_size=tile_size, tile_overlap=tile_overlap)

careamics/config/optimizer_models.py

@@ -1,6 +1,8 @@
+"""Optimizers and schedulers Pydantic models."""
+
 from __future__ import annotations
 
-from typing import Dict, Literal
+from typing import Literal
 
 from pydantic import (
     BaseModel,
@@ -19,8 +21,7 @@ from .support import SupportedOptimizer
 
 
 class OptimizerModel(BaseModel):
-    """
-    Torch optimizer.
+    """Torch optimizer Pydantic model.
 
     Only parameters supported by the corresponding torch optimizer will be taken
     into account. For more details, check:
@@ -31,7 +32,7 @@ class OptimizerModel(BaseModel):
 
     Attributes
     ----------
-    name : TorchOptimizer
+    name : {"Adam", "SGD"}
         Name of the optimizer.
     parameters : dict
         Parameters of the optimizer (see torch documentation).
@@ -55,7 +56,7 @@ class OptimizerModel(BaseModel):
 
     @field_validator("parameters")
     @classmethod
-    def filter_parameters(cls, user_params: dict, values: ValidationInfo) -> Dict:
+    def filter_parameters(cls, user_params: dict, values: ValidationInfo) -> dict:
         """
         Validate optimizer parameters.
 
@@ -70,7 +71,7 @@ class OptimizerModel(BaseModel):
 
         Returns
         -------
-        Dict
+        dict
             Filtered optimizer parameters.
 
         Raises
@@ -115,8 +116,7 @@ class OptimizerModel(BaseModel):
 
 
 class LrSchedulerModel(BaseModel):
-    """
-    Torch learning rate scheduler.
+    """Torch learning rate scheduler Pydantic model.
 
     Only parameters supported by the corresponding torch lr scheduler will be taken
     into account. For more details, check:
@@ -127,7 +127,7 @@ class LrSchedulerModel(BaseModel):
 
     Attributes
     ----------
-    name : TorchLRScheduler
+    name : {"ReduceLROnPlateau", "StepLR"}
         Name of the learning rate scheduler.
     parameters : dict
         Parameters of the learning rate scheduler (see torch documentation).
@@ -146,7 +146,7 @@ class LrSchedulerModel(BaseModel):
 
     @field_validator("parameters")
     @classmethod
-    def filter_parameters(cls, user_params: dict, values: ValidationInfo) -> Dict:
+    def filter_parameters(cls, user_params: dict, values: ValidationInfo) -> dict:
         """Filter parameters based on the learning rate scheduler's signature.
 
         Parameters
@@ -158,7 +158,7 @@ class LrSchedulerModel(BaseModel):
 
         Returns
         -------
-        Dict
+        dict
             Filtered scheduler parameters.
 
         Raises

careamics/config/support/__init__.py

@@ -14,7 +14,6 @@ __all__ = [
     "SupportedPixelManipulation",
     "SupportedTransform",
     "SupportedData",
-    "SupportedExtractionStrategy",
     "SupportedStructAxis",
     "SupportedLogger",
 ]
@@ -24,7 +23,6 @@ from .supported_activations import SupportedActivation
 from .supported_algorithms import SupportedAlgorithm
 from .supported_architectures import SupportedArchitecture
 from .supported_data import SupportedData
-from .supported_extraction_strategies import SupportedExtractionStrategy
 from .supported_loggers import SupportedLogger
 from .supported_losses import SupportedLoss
 from .supported_optimizers import SupportedOptimizer, SupportedScheduler

careamics/config/support/supported_activations.py

@@ -1,3 +1,5 @@
+"""Activations supported by CAREamics."""
+
 from careamics.utils import BaseEnum
 
 

careamics/config/support/supported_algorithms.py

@@ -1,3 +1,5 @@
+"""Algorithms supported by CAREamics."""
+
 from __future__ import annotations
 
 from careamics.utils import BaseEnum
@@ -10,9 +12,9 @@ class SupportedAlgorithm(str, BaseEnum):
     """
 
     N2V = "n2v"
-    CUSTOM = "custom"
     CARE = "care"
     N2N = "n2n"
+    CUSTOM = "custom"
     # PN2V = "pn2v"
     # HDN = "hdn"
     # SEG = "segmentation"

careamics/config/support/supported_architectures.py

@@ -1,3 +1,5 @@
+"""Architectures supported by CAREamics."""
+
 from careamics.utils import BaseEnum
 
 

careamics/config/support/supported_data.py

@@ -1,3 +1,5 @@
+"""Data supported by CAREamics."""
+
 from __future__ import annotations
 
 from typing import Union

careamics/config/support/supported_loggers.py

@@ -1,3 +1,5 @@
+"""Logger supported by CAREamics."""
+
 from careamics.utils import BaseEnum
 
 

careamics/config/support/supported_losses.py

@@ -1,3 +1,5 @@
+"""Losses supported by CAREamics."""
+
 from careamics.utils import BaseEnum
 
 

careamics/config/support/supported_optimizers.py

@@ -1,3 +1,5 @@
+"""Optimizers and schedulers supported by CAREamics."""
+
 from careamics.utils import BaseEnum
 
 

careamics/config/support/supported_pixel_manipulations.py

@@ -1,15 +1,15 @@
+"""Pixel manipulation methods supported by CAREamics."""
+
 from careamics.utils import BaseEnum
 
 
 class SupportedPixelManipulation(str, BaseEnum):
-    """_summary_.
+    """Supported Noise2Void pixel manipulations.
 
     - Uniform: Replace masked pixel value by a (uniformly) randomly selected neighbor
         pixel value.
     - Median: Replace masked pixel value by the mean of the neighborhood.
     """
 
-    # TODO docs
-
     UNIFORM = "uniform"
     MEDIAN = "median"

careamics/config/support/supported_struct_axis.py

@@ -1,3 +1,5 @@
+"""StructN2V axes supported by CAREamics."""
+
 from careamics.utils import BaseEnum