careamics 0.1.0rc2__py3-none-any.whl → 0.1.0rc3__py3-none-any.whl

careamics/config/inference_model.py

@@ -0,0 +1,283 @@
+"""Pydantic model representing CAREamics prediction configuration."""
+from __future__ import annotations
+
+from typing import Any, List, Literal, Optional, Union
+
+from albumentations import Compose
+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 InferenceModel(BaseModel):
+    """Configuration class for the prediction model."""
+
+    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(
+        default=None, min_length=2, max_length=3
+    )
+    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: Union[List[TRANSFORMS_UNION], Compose] = Field(
+        default=[
+            {
+                "name": SupportedTransform.NORMALIZE.value,
+            },
+        ],
+        validate_default=True,
+    )
+
+    # only default TTAs are supported for now
+    tta_transforms: bool = Field(default=True)
+
+    # Dataloader parameters
+    batch_size: int = Field(default=1, ge=1)
+
+    @field_validator("tile_overlap")
+    @classmethod
+    def all_elements_non_zero_even(
+        cls, patch_list: Optional[Union[List[int]]]
+    ) -> Optional[Union[List[int]]]:
+        """
+        Validate patch size.
+
+        Patch size must be non-zero, positive and even.
+
+        Parameters
+        ----------
+        patch_list : Optional[Union[List[int]]]
+            Patch size.
+
+        Returns
+        -------
+        Optional[Union[List[int]]]
+            Validated patch size.
+
+        Raises
+        ------
+        ValueError
+            If the patch size is 0.
+        ValueError
+            If the patch size is not even.
+        """
+        if patch_list is not None:
+            for dim in patch_list:
+                if dim < 1:
+                    raise ValueError(
+                        f"Patch size must be non-zero positive (got {dim})."
+                    )
+
+                if dim % 2 != 0:
+                    raise ValueError(f"Patch size must be even (got {dim}).")
+
+        return patch_list
+
+    @field_validator("tile_size")
+    @classmethod
+    def tile_min_8_power_of_2(
+        cls, tile_list: Optional[Union[List[int]]]
+    ) -> Optional[Union[List[int]]]:
+        """
+        Validate that each entry is greater or equal than 8 and a power of 2.
+
+        Parameters
+        ----------
+        tile_list : List[int]
+            Patch size.
+
+        Returns
+        -------
+        List[int]
+            Validated patch size.
+
+        Raises
+        ------
+        ValueError
+            If the patch size if smaller than 8.
+        ValueError
+            If the patch size is not a power of 2.
+        """
+        patch_size_ge_than_8_power_of_2(tile_list)
+
+        return tile_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_transforms(
+        cls, transforms: Union[List[TRANSFORMS_UNION], Compose]
+    ) -> Union[List[TRANSFORMS_UNION], Compose]:
+        """
+        Validate that transforms do not have N2V pixel manipulate transforms.
+
+        Parameters
+        ----------
+        transforms : Union[List[TransformModel], Compose]
+            Transforms.
+
+        Returns
+        -------
+        Union[List[Transformations_Union], Compose]
+            Validated transforms.
+
+        Raises
+        ------
+        ValueError
+            If transforms contain N2V pixel manipulate transforms.
+        """
+        if not isinstance(transforms, Compose) and 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:
+        """
+        Validate 2D/3D dimensions between axes and tile size.
+
+        Returns
+        -------
+        Self
+            Validated prediction model.
+        """
+        expected_len = 3 if "Z" in self.axes else 2
+
+        if self.tile_size is not None and self.tile_overlap is not None:
+            if len(self.tile_size) != expected_len:
+                raise ValueError(
+                    f"Tile size must have {expected_len} dimensions given axes "
+                    f"{self.axes} (got {self.tile_size})."
+                )
+
+            if len(self.tile_overlap) != expected_len:
+                raise ValueError(
+                    f"Tile overlap must have {expected_len} dimensions given axes "
+                    f"{self.axes} (got {self.tile_overlap})."
+                )
+
+            if any((i >= j) for i, j in zip(self.tile_overlap, self.tile_size)):
+                raise ValueError("Tile overlap must be smaller than tile size.")
+
+        return self
+
+    @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 prediction 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.mean is None) != (self.std is None):
+            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
+            if not isinstance(self.transforms, Compose):
+                for transform in self.transforms:
+                    if transform.name == SupportedTransform.NORMALIZE.value:
+                        transform.mean = self.mean
+                        transform.std = self.std
+
+        return self
+
+    def _update(self, **kwargs: Any) -> None:
+        """
+        Update multiple arguments at once.
+
+        Parameters
+        ----------
+        **kwargs : Any
+            Key-value pairs of arguments to update.
+        """
+        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:
+        """
+        Set 3D parameters.
+
+        Parameters
+        ----------
+        axes : str
+            Axes.
+        tile_size : List[int]
+            Tile size.
+        tile_overlap : List[int]
+            Tile overlap.
+        """
+        self._update(axes=axes, tile_size=tile_size, tile_overlap=tile_overlap)

careamics/config/noise_models.py

@@ -0,0 +1,162 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import Dict, Union
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+
+class NoiseModelType(str, Enum):
+    """
+    Available noise models.
+
+    Currently supported noise models:
+
+        - hist: Histogram noise model.
+        - gmm: Gaussian mixture model noise model.F
+    """
+
+    NONE = "none"
+    HIST = "hist"
+    GMM = "gmm"
+
+    # TODO add validator decorator
+    @classmethod
+    def validate_noise_model_type(
+        cls, noise_model: Union[str, NoiseModel], parameters: dict
+    ) -> None:
+        """_summary_.
+
+        Parameters
+        ----------
+        noise_model : Union[str, NoiseModel]
+            _description_
+        parameters : dict
+            _description_
+
+        Returns
+        -------
+        BaseModel
+            _description_
+        """
+        if noise_model == NoiseModelType.HIST.value:
+            HistogramNoiseModel(**parameters)
+            return HistogramNoiseModel().model_dump() if not parameters else parameters
+
+        elif noise_model == NoiseModelType.GMM.value:
+            GaussianMixtureNoiseModel(**parameters)
+            return (
+                GaussianMixtureNoiseModel().model_dump()
+                if not parameters
+                else parameters
+            )
+
+
+class NoiseModel(BaseModel):
+    """_summary_.
+
+    Parameters
+    ----------
+    BaseModel : _type_
+        _description_
+
+    Returns
+    -------
+    _type_
+        _description_
+
+    Raises
+    ------
+    ValueError
+        _description_
+    """
+
+    model_config = ConfigDict(
+        use_enum_values=True,
+        protected_namespaces=(),  # allows to use model_* as a field name
+        validate_assignment=True,
+    )
+
+    model_type: NoiseModelType
+    parameters: Dict = Field(default_factory=dict, validate_default=True)
+
+    @field_validator("parameters")
+    @classmethod
+    def validate_parameters(cls, data, values) -> Dict:
+        """_summary_.
+
+        Parameters
+        ----------
+        parameters : Dict
+            _description_
+
+        Returns
+        -------
+        Dict
+            _description_
+        """
+        if values.data["model_type"] not in [NoiseModelType.GMM, NoiseModelType.HIST]:
+            raise ValueError(
+                f"Incorrect noise model {values.data['model_type']}."
+                f"Please refer to the documentation"  # TODO add link to documentation
+            )
+
+        parameters = NoiseModelType.validate_noise_model_type(
+            values.data["model_type"], data
+        )
+        return parameters
+
+
+class HistogramNoiseModel(BaseModel):
+    """
+    Histogram noise model.
+
+    Attributes
+    ----------
+    min_value : float
+        Minimum value in the input.
+    max_value : float
+        Maximum value in the input.
+    bins : int
+        Number of bins of the histogram.
+    """
+
+    min_value: float = Field(default=350.0, ge=0.0, le=65535.0)
+    max_value: float = Field(default=6500.0, ge=0.0, le=65535.0)
+    bins: int = Field(default=256, ge=1)
+
+
+class GaussianMixtureNoiseModel(BaseModel):
+    """
+    Gaussian mixture model noise model.
+
+    Attributes
+    ----------
+    min_signal : float
+    Minimum signal intensity expected in the image.
+    max_signal : float
+    Maximum signal intensity expected in the image.
+    weight : array
+    A [3*n_gaussian, n_coeff] sized array containing the values of the weights
+    describing the noise model.
+    Each gaussian contributes three parameters (mean, standard deviation and weight),
+    hence the number of rows in `weight` are 3*n_gaussian.
+    If `weight = None`, the weight array is initialized using the `min_signal` and
+    `max_signal` parameters.
+    n_gaussian: int
+    Number of gaussians.
+    n_coeff: int
+    Number of coefficients to describe the functional relationship between gaussian
+    parameters and the signal.
+    2 implies a linear relationship, 3 implies a quadratic relationship and so on.
+    device: device
+    GPU device.
+    min_sigma: int
+    """
+
+    num_components: int = Field(default=3, ge=1)
+    min_value: float = Field(default=350.0, ge=0.0, le=65535.0)
+    max_value: float = Field(default=6500.0, ge=0.0, le=65535.0)
+    n_gaussian: int = Field(default=3, ge=1)
+    n_coeff: int = Field(default=2, ge=1)
+    min_sigma: int = Field(default=50, ge=1)

careamics/config/optimizer_models.py

@@ -0,0 +1,181 @@
+from __future__ import annotations
+
+from typing import Dict, Literal
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    ValidationInfo,
+    field_validator,
+    model_validator,
+)
+from torch import optim
+from typing_extensions import Self
+
+from careamics.utils.torch_utils import filter_parameters
+
+from .support import SupportedOptimizer
+
+
+class OptimizerModel(BaseModel):
+    """
+    Torch optimizer.
+
+    Only parameters supported by the corresponding torch optimizer will be taken
+    into account. For more details, check:
+    https://pytorch.org/docs/stable/optim.html#algorithms
+
+    Note that mandatory parameters (see the specific Optimizer signature in the
+    link above) must be provided. For example, SGD requires `lr`.
+
+    Attributes
+    ----------
+    name : TorchOptimizer
+        Name of the optimizer.
+    parameters : dict
+        Parameters of the optimizer (see torch documentation).
+    """
+
+    # Pydantic class configuration
+    model_config = ConfigDict(
+        validate_assignment=True,
+    )
+
+    # Mandatory field
+    name: Literal["Adam", "SGD"] = Field(default="Adam", validate_default=True)
+
+    # Optional parameters, empty dict default value to allow filtering dictionary
+    parameters: dict = Field(
+        default={
+            "lr": 1e-4,
+        },
+        validate_default=True,
+    )
+
+    @field_validator("parameters")
+    @classmethod
+    def filter_parameters(cls, user_params: dict, values: ValidationInfo) -> Dict:
+        """
+        Validate optimizer parameters.
+
+        This method filters out unknown parameters, given the optimizer name.
+
+        Parameters
+        ----------
+        user_params : dict
+            Parameters passed on to the torch optimizer.
+        values : ValidationInfo
+            Pydantic field validation info, used to get the optimizer name.
+
+        Returns
+        -------
+        Dict
+            Filtered optimizer parameters.
+
+        Raises
+        ------
+        ValueError
+            If the optimizer name is not specified.
+        """
+        optimizer_name = values.data["name"]
+
+        # retrieve the corresponding optimizer class
+        optimizer_class = getattr(optim, optimizer_name)
+
+        # filter the user parameters according to the optimizer's signature
+        parameters = filter_parameters(optimizer_class, user_params)
+
+        return parameters
+
+    @model_validator(mode="after")
+    def sgd_lr_parameter(self) -> Self:
+        """
+        Check that SGD optimizer has the mandatory `lr` parameter specified.
+
+        This is specific for PyTorch < 2.2.
+
+        Returns
+        -------
+        Self
+            Validated optimizer.
+
+        Raises
+        ------
+        ValueError
+            If the optimizer is SGD and the lr parameter is not specified.
+        """
+        if self.name == SupportedOptimizer.SGD and "lr" not in self.parameters:
+            raise ValueError(
+                "SGD optimizer requires `lr` parameter, check that it has correctly "
+                "been specified in `parameters`."
+            )
+
+        return self
+
+
+class LrSchedulerModel(BaseModel):
+    """
+    Torch learning rate scheduler.
+
+    Only parameters supported by the corresponding torch lr scheduler will be taken
+    into account. For more details, check:
+    https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate
+
+    Note that mandatory parameters (see the specific LrScheduler signature in the
+    link above) must be provided. For example, StepLR requires `step_size`.
+
+    Attributes
+    ----------
+    name : TorchLRScheduler
+        Name of the learning rate scheduler.
+    parameters : dict
+        Parameters of the learning rate scheduler (see torch documentation).
+    """
+
+    # Pydantic class configuration
+    model_config = ConfigDict(
+        validate_assignment=True,
+    )
+
+    # Mandatory field
+    name: Literal["ReduceLROnPlateau", "StepLR"] = Field(default="ReduceLROnPlateau")
+
+    # Optional parameters
+    parameters: dict = Field(default={}, validate_default=True)
+
+    @field_validator("parameters")
+    @classmethod
+    def filter_parameters(cls, user_params: dict, values: ValidationInfo) -> Dict:
+        """Filter parameters based on the learning rate scheduler's signature.
+
+        Parameters
+        ----------
+        user_params : dict
+            User parameters.
+        values : ValidationInfo
+            Pydantic field validation info, used to get the scheduler name.
+
+        Returns
+        -------
+        Dict
+            Filtered scheduler parameters.
+
+        Raises
+        ------
+        ValueError
+            If the scheduler is StepLR and the step_size parameter is not specified.
+        """
+        # retrieve the corresponding scheduler class
+        scheduler_class = getattr(optim.lr_scheduler, values.data["name"])
+
+        # filter the user parameters according to the scheduler's signature
+        parameters = filter_parameters(scheduler_class, user_params)
+
+        if values.data["name"] == "StepLR" and "step_size" not in parameters:
+            raise ValueError(
+                "StepLR scheduler requires `step_size` parameter, check that it has "
+                "correctly been specified in `parameters`."
+            )
+
+        return parameters

careamics/config/references/__init__.py

@@ -0,0 +1,45 @@
+"""Module containing references to the algorithm used in CAREamics."""
+
+__all__ = [
+    "N2V2Ref",
+    "N2VRef",
+    "StructN2VRef",
+    "N2VDescription",
+    "N2V2Description",
+    "StructN2VDescription",
+    "StructN2V2Description",
+    "N2V",
+    "N2V2",
+    "STRUCT_N2V",
+    "STRUCT_N2V2",
+    "CUSTOM",
+    "N2N",
+    "CARE",
+    "CAREDescription",
+    "N2NDescription",
+    "CARERef",
+    "N2NRef",
+]
+
+from .algorithm_descriptions import (
+    CARE,
+    CUSTOM,
+    N2N,
+    N2V,
+    N2V2,
+    STRUCT_N2V,
+    STRUCT_N2V2,
+    CAREDescription,
+    N2NDescription,
+    N2V2Description,
+    N2VDescription,
+    StructN2V2Description,
+    StructN2VDescription,
+)
+from .references import (
+    CARERef,
+    N2NRef,
+    N2V2Ref,
+    N2VRef,
+    StructN2VRef,
+)