careamics 0.1.0rc6__py3-none-any.whl → 0.1.0rc8__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,7 +18,6 @@ from typing_extensions import Annotated, Self
 
 from .support import SupportedTransform
 from .transformations.n2v_manipulate_model import N2VManipulateModel
-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
@@ -26,7 +26,6 @@ TRANSFORMS_UNION = Annotated[
     Union[
         XYFlipModel,
         XYRandomRotate90Model,
-        NormalizeModel,
         N2VManipulateModel,
     ],
     Discriminator("name"),  # used to tell the different transform models apart
@@ -39,7 +38,9 @@ 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.
 
@@ -54,8 +55,8 @@ class DataConfig(BaseModel):
     ...     axes="YX"
     ... )
 
-    To change the mean and std of the data:
-    >>> data.set_mean_and_std(mean=214.3, std=84.5)
+    To change the image_means and image_stds of the data:
+    >>> data.set_means_and_stds(image_means=[214.3], image_stds=[84.5])
 
     One can pass also a list of transformations, by keyword, using the
     SupportedTransform value:
@@ -67,11 +68,6 @@ class DataConfig(BaseModel):
     ...     axes="YX",
     ...     transforms=[
     ...         {
-    ...             "name": SupportedTransform.NORMALIZE.value,
-    ...             "mean": 167.6,
-    ...             "std": 47.2,
-    ...         },
-    ...         {
     ...             "name": "XYFlip",
     ...         }
     ...     ]
@@ -84,20 +80,41 @@ 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)
-    batch_size: int = Field(default=1, ge=1, validate_default=True)
+    data_type: Literal["array", "tiff", "custom"]
+    """Type of input data, numpy.ndarray (array) or paths (tiff and custom), as defined
+    in SupportedData."""
+
     axes: str
+    """Axes of the data, as defined in SupportedAxes."""
+
+    patch_size: Union[list[int]] = Field(..., min_length=2, max_length=3)
+    """Patch size, as used during training."""
+
+    batch_size: int = Field(default=1, ge=1, validate_default=True)
+    """Batch size for training."""
 
     # Optional fields
-    mean: Optional[float] = None
-    std: Optional[float] = None
+    image_means: Optional[list[float]] = Field(
+        default=None, min_length=0, max_length=32
+    )
+    """Means of the data across channels, used for normalization."""
+
+    image_stds: Optional[list[float]] = Field(default=None, min_length=0, max_length=32)
+    """Standard deviations of the data across channels, used for normalization."""
 
-    transforms: List[TRANSFORMS_UNION] = Field(
+    target_means: Optional[list[float]] = Field(
+        default=None, min_length=0, max_length=32
+    )
+    """Means of the target data across channels, used for normalization."""
+
+    target_stds: Optional[list[float]] = Field(
+        default=None, min_length=0, max_length=32
+    )
+    """Standard deviations of the target data across channels, used for
+    normalization."""
+
+    transforms: list[TRANSFORMS_UNION] = Field(
         default=[
-            {
-                "name": SupportedTransform.NORMALIZE.value,
-            },
             {
                 "name": SupportedTransform.XY_FLIP.value,
             },
@@ -110,14 +127,17 @@ class DataConfig(BaseModel):
         ],
         validate_default=True,
     )
+    """List of transformations to apply to the data, available transforms are defined
+    in SupportedTransform. The default values are set for Noise2Void."""
 
     dataloader_params: Optional[dict] = None
+    """Dictionary of PyTorch dataloader parameters."""
 
     @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.
 
@@ -125,12 +145,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
@@ -180,19 +200,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
@@ -235,29 +255,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 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:
-                    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
 
@@ -341,23 +365,48 @@ 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_means_and_stds(
+        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.
+        Set mean and standard deviation of the data across channels.
 
         This method should be used instead setting the fields directly, as it would
         otherwise trigger a validation error.
 
         Parameters
         ----------
-        mean : float
-            Mean of the data.
-        std : float
-            Standard deviation of the data.
+        image_means : numpy.ndarray ,tuple or list
+            Mean values for normalization.
+        image_stds : numpy.ndarray, tuple or list
+            Standard deviation values for normalization.
+        target_means : numpy.ndarray, tuple or list, optional
+            Target mean values for normalization, by default ().
+        target_stds : numpy.ndarray, tuple or list, optional
+            Target standard deviation values for normalization, by default ().
         """
-        self._update(mean=mean, 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.
 
@@ -365,7 +414,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,7 +2,7 @@
 
 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
@@ -15,31 +15,41 @@ class InferenceConfig(BaseModel):
 
     model_config = ConfigDict(validate_assignment=True, arbitrary_types_allowed=True)
 
-    # Mandatory fields
     data_type: Literal["array", "tiff", "custom"]  # As defined in SupportedData
-    tile_size: Optional[Union[List[int]]] = Field(
+    """Type of input data: numpy.ndarray (array) or path (tiff or custom)."""
+
+    tile_size: Optional[Union[list[int]]] = Field(
         default=None, min_length=2, max_length=3
     )
-    tile_overlap: Optional[Union[List[int]]] = Field(
+    """Tile size of prediction, only effective if `tile_overlap` is specified."""
+
+    tile_overlap: Optional[Union[list[int]]] = Field(
         default=None, min_length=2, max_length=3
     )
+    """Overlap between tiles, only effective if `tile_size` is specified."""
 
     axes: str
+    """Data axes (TSCZYX) in the order of the input data."""
 
-    mean: float
-    std: float = Field(..., ge=0.0)
+    image_means: list = Field(..., min_length=0, max_length=32)
+    """Mean values for each input channel."""
 
-    # only default TTAs are supported for now
+    image_stds: list = Field(..., min_length=0, max_length=32)
+    """Standard deviation values for each input channel."""
+
+    # TODO only default TTAs are supported for now
     tta_transforms: bool = Field(default=True)
+    """Whether to apply test-time augmentation (all 90 degrees rotations and flips)."""
 
     # Dataloader parameters
     batch_size: int = Field(default=1, ge=1)
+    """Batch size for prediction."""
 
     @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.
 
@@ -47,12 +57,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
@@ -77,19 +87,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
@@ -182,11 +192,23 @@ 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."
             )
 
+        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
 
     def _update(self, **kwargs: Any) -> None:
@@ -201,7 +223,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.
 
@@ -209,9 +231,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

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Dict, Literal
+from typing import Literal
 
 from pydantic import (
     BaseModel,
@@ -32,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).
@@ -45,6 +45,7 @@ class OptimizerModel(BaseModel):
 
     # Mandatory field
     name: Literal["Adam", "SGD"] = Field(default="Adam", validate_default=True)
+    """Name of the optimizer, supported optimizers are defined in SupportedOptimizer."""
 
     # Optional parameters, empty dict default value to allow filtering dictionary
     parameters: dict = Field(
@@ -53,10 +54,11 @@ class OptimizerModel(BaseModel):
         },
         validate_default=True,
     )
+    """Parameters of the optimizer, see PyTorch documentation for more details."""
 
     @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.
 
@@ -71,7 +73,7 @@ class OptimizerModel(BaseModel):
 
         Returns
         -------
-        Dict
+        dict
             Filtered optimizer parameters.
 
         Raises
@@ -127,7 +129,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).
@@ -140,13 +142,17 @@ class LrSchedulerModel(BaseModel):
 
     # Mandatory field
     name: Literal["ReduceLROnPlateau", "StepLR"] = Field(default="ReduceLROnPlateau")
+    """Name of the learning rate scheduler, supported schedulers are defined in
+    SupportedScheduler."""
 
     # Optional parameters
     parameters: dict = Field(default={}, validate_default=True)
+    """Parameters of the learning rate scheduler, see PyTorch documentation for more
+    details."""
 
     @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 +164,7 @@ class LrSchedulerModel(BaseModel):
 
         Returns
         -------
-        Dict
+        dict
             Filtered scheduler parameters.
 
         Raises

careamics/config/support/supported_data.py

@@ -60,9 +60,9 @@ class SupportedData(str, BaseEnum):
         return super()._missing_(value)
 
     @classmethod
-    def get_extension(cls, data_type: Union[str, SupportedData]) -> str:
+    def get_extension_pattern(cls, data_type: Union[str, SupportedData]) -> str:
         """
-        Path.rglob and fnmatch compatible extension.
+        Get Path.rglob and fnmatch compatible extension.
 
         Parameters
         ----------
@@ -72,13 +72,38 @@ class SupportedData(str, BaseEnum):
         Returns
         -------
         str
-            Corresponding extension.
+            Corresponding extension pattern.
         """
         if data_type == cls.ARRAY:
-            raise NotImplementedError(f"Data {data_type} are not loaded from file.")
+            raise NotImplementedError(f"Data '{data_type}' is not loaded from a file.")
         elif data_type == cls.TIFF:
             return "*.tif*"
         elif data_type == cls.CUSTOM:
             return "*.*"
         else:
             raise ValueError(f"Data type {data_type} is not supported.")
+
+    @classmethod
+    def get_extension(cls, data_type: Union[str, SupportedData]) -> str:
+        """
+        Get file extension of corresponding data type.
+
+        Parameters
+        ----------
+        data_type : str or SupportedData
+            Data type.
+
+        Returns
+        -------
+        str
+            Corresponding extension.
+        """
+        if data_type == cls.ARRAY:
+            raise NotImplementedError(f"Data '{data_type}' is not loaded from a file.")
+        elif data_type == cls.TIFF:
+            return ".tiff"
+        elif data_type == cls.CUSTOM:
+            # TODO: improve this message
+            raise NotImplementedError("Custom extensions have to be passed elsewhere.")
+        else:
+            raise ValueError(f"Data type {data_type} is not supported.")

careamics/config/support/supported_transforms.py

@@ -8,5 +8,4 @@ class SupportedTransform(str, BaseEnum):
 
     XY_FLIP = "XYFlip"
     XY_RANDOM_ROTATE90 = "XYRandomRotate90"
-    NORMALIZE = "Normalize"
     N2V_MANIPULATE = "N2VManipulate"