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

careamics/config/inference_model.py

@@ -0,0 +1,239 @@
+"""Pydantic model representing CAREamics prediction configuration."""
+
+from __future__ import annotations
+
+from typing import Any, Literal, Optional, Union
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
+from typing_extensions import Self
+
+from .validators import check_axes_validity, patch_size_ge_than_8_power_of_2
+
+
+class InferenceConfig(BaseModel):
+    """Configuration class for the prediction model."""
+
+    model_config = ConfigDict(validate_assignment=True, arbitrary_types_allowed=True)
+
+    data_type: Literal["array", "tiff", "custom"]  # As defined in SupportedData
+    """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 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."""
+
+    image_means: list = Field(..., min_length=0, max_length=32)
+    """Mean values for each input channel."""
+
+    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[list[int]]
+    ) -> Optional[list[int]]:
+        """
+        Validate tile overlap.
+
+        Overlaps must be non-zero, positive and even.
+
+        Parameters
+        ----------
+        tile_overlap : list[int] or None
+            Patch size.
+
+        Returns
+        -------
+        list[int] or None
+            Validated tile overlap.
+
+        Raises
+        ------
+        ValueError
+            If the patch size is 0.
+        ValueError
+            If the patch size is not even.
+        """
+        if tile_overlap is not None:
+            for dim in tile_overlap:
+                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 tile_overlap
+
+    @field_validator("tile_size")
+    @classmethod
+    def tile_min_8_power_of_2(
+        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 of int
+            Patch size.
+
+        Returns
+        -------
+        list of 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
+
+    @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 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:
+        """
+        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 of int
+            Tile size.
+        tile_overlap : list of int
+            Tile overlap.
+        """
+        self._update(axes=axes, tile_size=tile_size, tile_overlap=tile_overlap)

careamics/config/optimizer_models.py

@@ -0,0 +1,187 @@
+"""Optimizers and schedulers Pydantic models."""
+
+from __future__ import annotations
+
+from typing import 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 Pydantic model.
+
+    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 : {"Adam", "SGD"}
+        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)
+    """Name of the optimizer, supported optimizers are defined in SupportedOptimizer."""
+
+    # Optional parameters, empty dict default value to allow filtering dictionary
+    parameters: dict = Field(
+        default={
+            "lr": 1e-4,
+        },
+        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:
+        """
+        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 Pydantic model.
+
+    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 : {"ReduceLROnPlateau", "StepLR"}
+        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")
+    """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:
+        """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,
+)

careamics/config/references/algorithm_descriptions.py

@@ -0,0 +1,132 @@
+"""Descriptions of the algorithms used in CAREmics."""
+
+from pydantic import BaseModel
+
+CUSTOM = "Custom"
+N2V = "Noise2Void"
+N2V2 = "N2V2"
+STRUCT_N2V = "StructN2V"
+STRUCT_N2V2 = "StructN2V2"
+N2N = "Noise2Noise"
+CARE = "CARE"
+
+
+N2V_DESCRIPTION = (
+    "Noise2Void is a UNet-based self-supervised algorithm that "
+    "uses blind-spot training to denoise images. In short, in every "
+    "patches during training, random pixels are selected and their "
+    "value replaced by a neighboring pixel value. The network is then "
+    "trained to predict the original pixel value. The algorithm "
+    "relies on the continuity of the signal (neighboring pixels have "
+    "similar values) and the pixel-wise independence of the noise "
+    "(the noise in a pixel is not correlated with the noise in "
+    "neighboring pixels)."
+)
+
+
+class AlgorithmDescription(BaseModel):
+    """Description of an algorithm.
+
+    Attributes
+    ----------
+    description : str
+        Description of the algorithm.
+    """
+
+    description: str
+
+
+class N2VDescription(AlgorithmDescription):
+    """Description of Noise2Void.
+
+    Attributes
+    ----------
+    description : str
+        Description of Noise2Void.
+    """
+
+    description: str = N2V_DESCRIPTION
+
+
+class N2V2Description(AlgorithmDescription):
+    """Description of N2V2.
+
+    Attributes
+    ----------
+    description : str
+        Description of N2V2.
+    """
+
+    description: str = (
+        "N2V2 is a variant of Noise2Void. "
+        + N2V_DESCRIPTION
+        + "\nN2V2 introduces blur-pool layers and removed skip "
+        "connections in the UNet architecture to remove checkboard "
+        "artefacts, a common artefacts ocurring in Noise2Void."
+    )
+
+
+class StructN2VDescription(AlgorithmDescription):
+    """Description of StructN2V.
+
+    Attributes
+    ----------
+    description : str
+        Description of StructN2V.
+    """
+
+    description: str = (
+        "StructN2V is a variant of Noise2Void. "
+        + N2V_DESCRIPTION
+        + "\nStructN2V uses a linear mask (horizontal or vertical) to replace "
+        "the pixel values of neighbors of the masked pixels by a random "
+        "value. Such masking allows removing 1D structured noise from the "
+        "the images, the main failure case of the original N2V."
+    )
+
+
+class StructN2V2Description(AlgorithmDescription):
+    """Description of StructN2V2.
+
+    Attributes
+    ----------
+    description : str
+        Description of StructN2V2.
+    """
+
+    description: str = (
+        "StructN2V2 is a a variant of Noise2Void that uses both "
+        "structN2V and N2V2. "
+        + N2V_DESCRIPTION
+        + "\nStructN2V2 uses a linear mask (horizontal or vertical) to replace "
+        "the pixel values of neighbors of the masked pixels by a random "
+        "value. Such masking allows removing 1D structured noise from the "
+        "the images, the main failure case of the original N2V."
+        "\nN2V2 introduces blur-pool layers and removed skip connections in "
+        "the UNet architecture to remove checkboard artefacts, a common "
+        "artefacts ocurring in Noise2Void."
+    )
+
+
+class N2NDescription(AlgorithmDescription):
+    """Description of Noise2Noise.
+
+    Attributes
+    ----------
+    description : str
+        Description of Noise2Noise.
+    """
+
+    description: str = "Noise2Noise"  # TODO
+
+
+class CAREDescription(AlgorithmDescription):
+    """Description of CARE.
+
+    Attributes
+    ----------
+    description : str
+        Description of CARE.
+    """
+
+    description: str = "CARE"  # TODO

careamics/config/references/references.py

@@ -0,0 +1,39 @@
+"""References for the CAREamics algorithms."""
+
+from bioimageio.spec.generic.v0_3 import CiteEntry
+
+N2VRef = CiteEntry(
+    text='Krull, A., Buchholz, T.O. and Jug, F., 2019. "Noise2Void - Learning '
+    'denoising from single noisy images". In Proceedings of the IEEE/CVF '
+    "conference on computer vision and pattern recognition (pp. 2129-2137).",
+    doi="10.1109/cvpr.2019.00223",
+)
+
+N2V2Ref = CiteEntry(
+    text="Höck, E., Buchholz, T.O., Brachmann, A., Jug, F. and Freytag, A., "
+    '2022. "N2V2 - Fixing Noise2Void checkerboard artifacts with modified '
+    'sampling strategies and a tweaked network architecture". In European '
+    "Conference on Computer Vision (pp. 503-518).",
+    doi="10.1007/978-3-031-25069-9_33",
+)
+
+StructN2VRef = CiteEntry(
+    text="Broaddus, C., Krull, A., Weigert, M., Schmidt, U. and Myers, G., 2020."
+    '"Removing structured noise with self-supervised blind-spot '
+    'networks". In 2020 IEEE 17th International Symposium on Biomedical '
+    "Imaging (ISBI) (pp. 159-163).",
+    doi="10.1109/isbi45749.2020.9098336",
+)
+
+N2NRef = CiteEntry(
+    text="Lehtinen, J., Munkberg, J., Hasselgren, J., Laine, S., Karras, T., "
+    'Aittala, M. and Aila, T., 2018. "Noise2Noise: Learning image restoration '
+    'without clean data". arXiv preprint arXiv:1803.04189.',
+    doi="10.48550/arXiv.1803.04189",
+)
+
+CARERef = CiteEntry(
+    text='Weigert, Martin, et al. "Content-aware image restoration: pushing the '
+    'limits of fluorescence microscopy." Nature methods 15.12 (2018): 1090-1097.',
+    doi="10.1038/s41592-018-0216-7",
+)

careamics/config/support/__init__.py

@@ -0,0 +1,31 @@
+"""Supported configuration options.
+
+Used throughout the code to ensure consistency. These should be kept in sync with the
+corresponding configuration options in the Pydantic models.
+"""
+
+__all__ = [
+    "SupportedArchitecture",
+    "SupportedActivation",
+    "SupportedOptimizer",
+    "SupportedScheduler",
+    "SupportedLoss",
+    "SupportedAlgorithm",
+    "SupportedPixelManipulation",
+    "SupportedTransform",
+    "SupportedData",
+    "SupportedStructAxis",
+    "SupportedLogger",
+]
+
+
+from .supported_activations import SupportedActivation
+from .supported_algorithms import SupportedAlgorithm
+from .supported_architectures import SupportedArchitecture
+from .supported_data import SupportedData
+from .supported_loggers import SupportedLogger
+from .supported_losses import SupportedLoss
+from .supported_optimizers import SupportedOptimizer, SupportedScheduler
+from .supported_pixel_manipulations import SupportedPixelManipulation
+from .supported_struct_axis import SupportedStructAxis
+from .supported_transforms import SupportedTransform

careamics/config/support/supported_activations.py

@@ -0,0 +1,26 @@
+"""Activations supported by CAREamics."""
+
+from careamics.utils import BaseEnum
+
+
+class SupportedActivation(str, BaseEnum):
+    """Supported activation functions.
+
+    - None, no activation will be used.
+    - Sigmoid
+    - Softmax
+    - Tanh
+    - ReLU
+    - LeakyReLU
+
+    All activations are defined in PyTorch.
+
+    See: https://pytorch.org/docs/stable/nn.html#loss-functions
+    """
+
+    NONE = "None"
+    SIGMOID = "Sigmoid"
+    SOFTMAX = "Softmax"
+    TANH = "Tanh"
+    RELU = "ReLU"
+    LEAKYRELU = "LeakyReLU"