careamics 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

careamics/config/data_model.py

@@ -0,0 +1,527 @@
+"""Data configuration."""
+
+from __future__ import annotations
+
+from pprint import pformat
+from typing import Any, Literal, Optional, Union
+
+import numpy as np
+from numpy.typing import NDArray
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Discriminator,
+    Field,
+    PlainSerializer,
+    field_validator,
+    model_validator,
+)
+from typing_extensions import Annotated, Self
+
+from .support import SupportedTransform
+from .transformations.n2v_manipulate_model import N2VManipulateModel
+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
+
+
+def np_float_to_scientific_str(x: float) -> str:
+    """Return a string scientific representation of a float.
+
+    In particular, this method is used to serialize floats to strings, allowing
+    numpy.float32 to be passed in the Pydantic model and written to a yaml file as str.
+
+    Parameters
+    ----------
+    x : float
+        Input value.
+
+    Returns
+    -------
+    str
+        Scientific string representation of the input value.
+    """
+    return np.format_float_scientific(x, precision=7)
+
+
+Float = Annotated[float, PlainSerializer(np_float_to_scientific_str, return_type=str)]
+"""Annotated float type, used to serialize floats to strings."""
+
+
+TRANSFORMS_UNION = Annotated[
+    Union[
+        XYFlipModel,
+        XYRandomRotate90Model,
+        N2VManipulateModel,
+    ],
+    Discriminator("name"),  # used to tell the different transform models apart
+]
+"""Available transforms in CAREamics."""
+
+
+class DataConfig(BaseModel):
+    """
+    Data configuration.
+
+    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. 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
+    --------
+    Minimum example:
+
+    >>> data = DataConfig(
+    ...     data_type="array", # defined in SupportedData
+    ...     patch_size=[128, 128],
+    ...     batch_size=4,
+    ...     axes="YX"
+    ... )
+
+    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:
+    >>> from careamics.config.support import SupportedTransform
+    >>> data = DataConfig(
+    ...     data_type="tiff",
+    ...     patch_size=[128, 128],
+    ...     batch_size=4,
+    ...     axes="YX",
+    ...     transforms=[
+    ...         {
+    ...             "name": "XYFlip",
+    ...         }
+    ...     ]
+    ... )
+    """
+
+    # Pydantic class configuration
+    model_config = ConfigDict(
+        validate_assignment=True,
+    )
+
+    # Dataset configuration
+    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
+    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."""
+
+    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.XY_FLIP.value,
+            },
+            {
+                "name": SupportedTransform.XY_RANDOM_ROTATE90.value,
+            },
+            {
+                "name": SupportedTransform.N2V_MANIPULATE.value,
+            },
+        ],
+        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]]:
+        """
+        Validate patch size.
+
+        Patch size must be powers of 2 and minimum 8.
+
+        Parameters
+        ----------
+        patch_list : list of int
+            Patch size.
+
+        Returns
+        -------
+        list of int
+            Validated patch size.
+
+        Raises
+        ------
+        ValueError
+            If the patch size is smaller than 8.
+        ValueError
+            If the patch size is not a power of 2.
+        """
+        patch_size_ge_than_8_power_of_2(patch_list)
+
+        return patch_list
+
+    @field_validator("axes")
+    @classmethod
+    def axes_valid(cls, axes: str) -> str:
+        """
+        Validate axes.
+
+        Axes must:
+        - be a combination of 'STCZYX'
+        - not contain duplicates
+        - contain at least 2 contiguous axes: X and Y
+        - contain at most 4 axes
+        - not contain both S and T axes
+
+        Parameters
+        ----------
+        axes : str
+            Axes to validate.
+
+        Returns
+        -------
+        str
+            Validated axes.
+
+        Raises
+        ------
+        ValueError
+            If axes are not valid.
+        """
+        # Validate axes
+        check_axes_validity(axes)
+
+        return axes
+
+    @field_validator("transforms")
+    @classmethod
+    def validate_prediction_transforms(
+        cls, transforms: list[TRANSFORMS_UNION]
+    ) -> list[TRANSFORMS_UNION]:
+        """
+        Validate N2VManipulate transform position in the transform list.
+
+        Parameters
+        ----------
+        transforms : list[Transformations_Union]
+            Transforms.
+
+        Returns
+        -------
+        list of transforms
+            Validated transforms.
+
+        Raises
+        ------
+        ValueError
+            If multiple instances of N2VManipulate are found.
+        """
+        transform_list = [t.name for t in transforms]
+
+        if SupportedTransform.N2V_MANIPULATE in transform_list:
+            # multiple N2V_MANIPULATE
+            if transform_list.count(SupportedTransform.N2V_MANIPULATE.value) > 1:
+                raise ValueError(
+                    f"Multiple instances of "
+                    f"{SupportedTransform.N2V_MANIPULATE} transforms "
+                    f"are not allowed."
+                )
+
+            # N2V_MANIPULATE not the last transform
+            elif transform_list[-1] != SupportedTransform.N2V_MANIPULATE:
+                index = transform_list.index(SupportedTransform.N2V_MANIPULATE.value)
+                transform = transforms.pop(index)
+                transforms.append(transform)
+
+        return transforms
+
+    @model_validator(mode="after")
+    def std_only_with_mean(self: Self) -> Self:
+        """
+        Check that mean and std are either both None, or both specified.
+
+        Returns
+        -------
+        Self
+            Validated data model.
+
+        Raises
+        ------
+        ValueError
+            If std is not None and mean is None.
+        """
+        # check that mean and std are either both None, or both specified
+        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.")
+
+        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 "
+            )
+
+        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
+
+    @model_validator(mode="after")
+    def validate_dimensions(self: Self) -> Self:
+        """
+        Validate 2D/3D dimensions between axes, patch size and transforms.
+
+        Returns
+        -------
+        Self
+            Validated data model.
+
+        Raises
+        ------
+        ValueError
+            If the transforms are not valid.
+        """
+        if "Z" in self.axes:
+            if len(self.patch_size) != 3:
+                raise ValueError(
+                    f"Patch size must have 3 dimensions if the data is 3D "
+                    f"({self.axes})."
+                )
+
+        else:
+            if len(self.patch_size) != 2:
+                raise ValueError(
+                    f"Patch size must have 3 dimensions if the data is 3D "
+                    f"({self.axes})."
+                )
+
+        return self
+
+    def __str__(self) -> str:
+        """
+        Pretty string reprensenting the configuration.
+
+        Returns
+        -------
+        str
+            Pretty string.
+        """
+        return pformat(self.model_dump())
+
+    def _update(self, **kwargs: Any) -> None:
+        """
+        Update multiple arguments at once.
+
+        Parameters
+        ----------
+        **kwargs : Any
+            Keyword arguments to update.
+        """
+        self.__dict__.update(kwargs)
+        self.__class__.model_validate(self.__dict__)
+
+    def has_n2v_manipulate(self) -> bool:
+        """
+        Check if the transforms contain N2VManipulate.
+
+        Returns
+        -------
+        bool
+            True if the transforms contain N2VManipulate, False otherwise.
+        """
+        return any(
+            transform.name == SupportedTransform.N2V_MANIPULATE.value
+            for transform in self.transforms
+        )
+
+    def add_n2v_manipulate(self) -> None:
+        """Add N2VManipulate to the transforms."""
+        if not self.has_n2v_manipulate():
+            self.transforms.append(
+                N2VManipulateModel(name=SupportedTransform.N2V_MANIPULATE.value)
+            )
+
+    def remove_n2v_manipulate(self) -> None:
+        """Remove N2VManipulate from the transforms."""
+        if self.has_n2v_manipulate():
+            self.transforms.pop(-1)
+
+    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 across channels.
+
+        This method should be used instead setting the fields directly, as it would
+        otherwise trigger a validation error.
+
+        Parameters
+        ----------
+        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 ().
+        """
+        # 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:
+        """
+        Set 3D parameters.
+
+        Parameters
+        ----------
+        axes : str
+            Axes.
+        patch_size : list of int
+            Patch size.
+        """
+        self._update(axes=axes, patch_size=patch_size)
+
+    def set_N2V2(self, use_n2v2: bool) -> None:
+        """
+        Set N2V2.
+
+        Parameters
+        ----------
+        use_n2v2 : bool
+            Whether to use N2V2.
+
+        Raises
+        ------
+        ValueError
+            If the N2V pixel manipulate transform is not found in the transforms.
+        """
+        if use_n2v2:
+            self.set_N2V2_strategy("median")
+        else:
+            self.set_N2V2_strategy("uniform")
+
+    def set_N2V2_strategy(self, strategy: Literal["uniform", "median"]) -> None:
+        """
+        Set N2V2 strategy.
+
+        Parameters
+        ----------
+        strategy : Literal["uniform", "median"]
+            Strategy to use for N2V2.
+
+        Raises
+        ------
+        ValueError
+            If the N2V pixel manipulate transform is not found in the transforms.
+        """
+        found_n2v = False
+
+        for transform in self.transforms:
+            if transform.name == SupportedTransform.N2V_MANIPULATE.value:
+                transform.strategy = strategy
+                found_n2v = True
+
+        if not found_n2v:
+            transforms = [t.name for t in self.transforms]
+            raise ValueError(
+                f"N2V_Manipulate transform not found in the transforms list "
+                f"({transforms})."
+            )
+
+    def set_structN2V_mask(
+        self, mask_axis: Literal["horizontal", "vertical", "none"], mask_span: int
+    ) -> None:
+        """
+        Set structN2V mask parameters.
+
+        Setting `mask_axis` to `none` will disable structN2V.
+
+        Parameters
+        ----------
+        mask_axis : Literal["horizontal", "vertical", "none"]
+            Axis along which to apply the mask. `none` will disable structN2V.
+        mask_span : int
+            Total span of the mask in pixels.
+
+        Raises
+        ------
+        ValueError
+            If the N2V pixel manipulate transform is not found in the transforms.
+        """
+        found_n2v = False
+
+        for transform in self.transforms:
+            if transform.name == SupportedTransform.N2V_MANIPULATE.value:
+                transform.struct_mask_axis = mask_axis
+                transform.struct_mask_span = mask_span
+                found_n2v = True
+
+        if not found_n2v:
+            transforms = [t.name for t in self.transforms]
+            raise ValueError(
+                f"N2V pixel manipulate transform not found in the transforms "
+                f"({transforms})."
+            )

careamics/config/fcn_algorithm_model.py

@@ -0,0 +1,147 @@
+"""Module containing `FCNAlgorithmConfig` class."""
+
+from pprint import pformat
+from typing import Literal, Union
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+from typing_extensions import Self
+
+from careamics.config.architectures import CustomModel, UNetModel
+from careamics.config.optimizer_models import LrSchedulerModel, OptimizerModel
+
+
+class FCNAlgorithmConfig(BaseModel):
+    """Algorithm configuration.
+
+    This Pydantic model validates the parameters governing the components of the
+    training algorithm: which algorithm, loss function, model architecture, optimizer,
+    and learning rate scheduler to use.
+
+    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
+    ----------
+    algorithm : Literal["n2v", "custom"]
+        Algorithm to use.
+    loss : Literal["n2v", "mae", "mse"]
+        Loss function to use.
+    model : Union[UNetModel, LVAEModel, CustomModel]
+        Model architecture to use.
+    optimizer : OptimizerModel, optional
+        Optimizer to use.
+    lr_scheduler : LrSchedulerModel, optional
+        Learning rate scheduler to use.
+
+    Raises
+    ------
+    ValueError
+        Algorithm parameter type validation errors.
+    ValueError
+        If the algorithm, loss and model are not compatible.
+
+    Examples
+    --------
+    Minimum example:
+    >>> from careamics.config import FCNAlgorithmConfig
+    >>> config_dict = {
+    ...     "algorithm": "n2v",
+    ...     "algorithm_type": "fcn",
+    ...     "loss": "n2v",
+    ...     "model": {
+    ...         "architecture": "UNet",
+    ...     }
+    ... }
+    >>> config = FCNAlgorithmConfig(**config_dict)
+    """
+
+    # Pydantic class configuration
+    model_config = ConfigDict(
+        protected_namespaces=(),  # allows to use model_* as a field name
+        validate_assignment=True,
+        extra="allow",
+    )
+
+    # Mandatory fields
+    # defined in SupportedAlgorithm
+    algorithm_type: Literal["fcn"]
+    """Algorithm type must be `fcn` (fully convolutional network) to differentiate this
+    configuration from LVAE."""
+
+    algorithm: Literal["n2v", "care", "n2n", "custom"]
+    """Name of the algorithm, as defined in SupportedAlgorithm. Use `custom` for custom
+    model architecture."""
+
+    loss: Literal["n2v", "mae", "mse"]
+    """Loss function to use, as defined in SupportedLoss."""
+
+    model: Union[UNetModel, CustomModel] = Field(discriminator="architecture")
+    """Model architecture to use, along with its parameters. Compatible architectures
+    are defined in SupportedArchitecture, and their Pydantic models in
+    `careamics.config.architectures`."""
+    # TODO supported architectures are now all the architectures but does not warn users
+    # of the compatibility with the algorithm
+
+    # Optional fields
+    optimizer: OptimizerModel = OptimizerModel()
+    """Optimizer to use, defined in SupportedOptimizer."""
+
+    lr_scheduler: LrSchedulerModel = LrSchedulerModel()
+    """Learning rate scheduler to use, defined in SupportedLrScheduler."""
+
+    @model_validator(mode="after")
+    def algorithm_cross_validation(self: Self) -> Self:
+        """Validate the algorithm model based on `algorithm`.
+
+        N2V:
+        - loss must be n2v
+        - model must be a `UNetModel`
+
+        Returns
+        -------
+        Self
+            The validated model.
+        """
+        # N2V
+        if self.algorithm == "n2v":
+            # n2v is only compatible with the n2v loss
+            if self.loss != "n2v":
+                raise ValueError(
+                    f"Algorithm {self.algorithm} only supports loss `n2v`."
+                )
+
+            # n2v 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`."
+                )
+
+            # n2v requires the number of input and output channels to be the same
+            if self.model.in_channels != self.model.num_classes:
+                raise ValueError(
+                    "N2V 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`.")
+
+        if (self.algorithm == "custom") != (self.model.architecture == "custom"):
+            raise ValueError(
+                "Algorithm and model architecture must be both `custom` or not."
+            )
+
+        return self
+
+    def __str__(self) -> str:
+        """Pretty string representing the configuration.
+
+        Returns
+        -------
+        str
+            Pretty string.
+        """
+        return pformat(self.model_dump())