careamics 0.0.5__py3-none-any.whl → 0.0.7__py3-none-any.whl

careamics/config/{data_model.py → data/data_model.py}

@@ -2,8 +2,10 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 from pprint import pformat
-from typing import Any, Literal, Optional, Union
+from typing import Annotated, Any, Literal, Optional, Union
+from warnings import warn
 
 import numpy as np
 from numpy.typing import NDArray
@@ -15,11 +17,10 @@ from pydantic import (
     field_validator,
     model_validator,
 )
-from typing_extensions import Annotated, Self
+from typing_extensions import Self
 
-from .support import SupportedTransform
-from .transformations import TRANSFORMS_UNION, N2VManipulateModel
-from .validators import check_axes_validity, patch_size_ge_than_8_power_of_2
+from ..transformations import N2V_TRANSFORMS_UNION, XYFlipModel, XYRandomRotate90Model
+from ..validators import check_axes_validity, patch_size_ge_than_8_power_of_2
 
 
 def np_float_to_scientific_str(x: float) -> str:
@@ -45,47 +46,8 @@ Float = Annotated[float, PlainSerializer(np_float_to_scientific_str, return_type
 """Annotated float type, used to serialize floats to strings."""
 
 
-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",
-    ...         }
-    ...     ]
-    ... )
-    """
+class GeneralDataConfig(BaseModel):
+    """General data configuration."""
 
     # Pydantic class configuration
     model_config = ConfigDict(
@@ -126,25 +88,26 @@ class DataConfig(BaseModel):
     """Standard deviations of the target data across channels, used for
     normalization."""
 
-    transforms: list[TRANSFORMS_UNION] = Field(
+    # defining as Sequence allows assigning subclasses of TransformModel without mypy
+    # complaining, this is important for instance to differentiate N2VDataConfig and
+    # DataConfig
+    transforms: Sequence[N2V_TRANSFORMS_UNION] = Field(
         default=[
-            {
-                "name": SupportedTransform.XY_FLIP.value,
-            },
-            {
-                "name": SupportedTransform.XY_RANDOM_ROTATE90.value,
-            },
-            {
-                "name": SupportedTransform.N2V_MANIPULATE.value,
-            },
+            XYFlipModel(),
+            XYRandomRotate90Model(),
         ],
         validate_default=True,
     )
     """List of transformations to apply to the data, available transforms are defined
-    in SupportedTransform. The default values are set for Noise2Void."""
+    in SupportedTransform."""
+
+    train_dataloader_params: dict[str, Any] = Field(
+        default={"shuffle": True}, validate_default=True
+    )
+    """Dictionary of PyTorch training dataloader parameters."""
 
-    dataloader_params: Optional[dict] = None
-    """Dictionary of PyTorch dataloader parameters."""
+    val_dataloader_params: dict[str, Any] = Field(default={})
+    """Dictionary of PyTorch validation dataloader parameters."""
 
     @field_validator("patch_size")
     @classmethod
@@ -210,47 +173,44 @@ class DataConfig(BaseModel):
 
         return axes
 
-    @field_validator("transforms")
+    @field_validator("train_dataloader_params")
     @classmethod
-    def validate_prediction_transforms(
-        cls, transforms: list[TRANSFORMS_UNION]
-    ) -> list[TRANSFORMS_UNION]:
+    def shuffle_train_dataloader(
+        cls, train_dataloader_params: dict[str, Any]
+    ) -> dict[str, Any]:
         """
-        Validate N2VManipulate transform position in the transform list.
+        Validate that "shuffle" is included in the training dataloader params.
+
+        A warning will be raised if `shuffle=False`.
 
         Parameters
         ----------
-        transforms : list[Transformations_Union]
-            Transforms.
+        train_dataloader_params : dict of {str: Any}
+            The training dataloader parameters.
 
         Returns
         -------
-        list of transforms
-            Validated transforms.
+        dict of {str: Any}
+            The validated training dataloader parameters.
 
         Raises
         ------
         ValueError
-            If multiple instances of N2VManipulate are found.
+            If "shuffle" is not included in the training dataloader params.
         """
-        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
+        if "shuffle" not in train_dataloader_params:
+            raise ValueError(
+                "Value for 'shuffle' was not included in the `train_dataloader_params`."
+            )
+        elif ("shuffle" in train_dataloader_params) and (
+            not train_dataloader_params["shuffle"]
+        ):
+            warn(
+                "Dataloader parameters include `shuffle=False`, this will be passed to "
+                "the training dataloader and may result in bad results.",
+                stacklevel=1,
+            )
+        return train_dataloader_params
 
     @model_validator(mode="after")
     def std_only_with_mean(self: Self) -> Self:
@@ -350,32 +310,6 @@ class DataConfig(BaseModel):
         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],
@@ -430,84 +364,55 @@ class DataConfig(BaseModel):
         """
         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
+class DataConfig(GeneralDataConfig):
+    """
+    Data configuration.
 
-        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})."
-            )
+    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.
 
-    def set_structN2V_mask(
-        self, mask_axis: Literal["horizontal", "vertical", "none"], mask_span: int
-    ) -> None:
-        """
-        Set structN2V mask parameters.
+    All supported transforms are defined in the SupportedTransform enum.
 
-        Setting `mask_axis` to `none` will disable structN2V.
+    Examples
+    --------
+    Minimum example:
 
-        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.
+    >>> data = DataConfig(
+    ...     data_type="array", # defined in SupportedData
+    ...     patch_size=[128, 128],
+    ...     batch_size=4,
+    ...     axes="YX"
+    ... )
 
-        Raises
-        ------
-        ValueError
-            If the N2V pixel manipulate transform is not found in the transforms.
-        """
-        found_n2v = False
+    To change the image_means and image_stds of the data:
+    >>> data.set_means_and_stds(image_means=[214.3], image_stds=[84.5])
 
-        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
+    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",
+    ...         }
+    ...     ]
+    ... )
+    """
 
-        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})."
-            )
+    transforms: Sequence[Union[XYFlipModel, XYRandomRotate90Model]] = Field(
+        default=[
+            XYFlipModel(),
+            XYRandomRotate90Model(),
+        ],
+        validate_default=True,
+    )
+    """List of transformations to apply to the data, available transforms are defined
+    in SupportedTransform. This excludes N2V specific transformations."""

careamics/config/data/n2v_data_model.py

@@ -0,0 +1,193 @@
+"""Noise2Void specific data configuration model."""
+
+from collections.abc import Sequence
+from typing import Literal
+
+from pydantic import Field, field_validator
+
+from careamics.config.data.data_model import GeneralDataConfig
+from careamics.config.support import SupportedTransform
+from careamics.config.transformations import (
+    N2V_TRANSFORMS_UNION,
+    N2VManipulateModel,
+    XYFlipModel,
+    XYRandomRotate90Model,
+)
+
+
+class N2VDataConfig(GeneralDataConfig):
+    """N2V specific data configuration model."""
+
+    transforms: Sequence[N2V_TRANSFORMS_UNION] = Field(
+        default=[XYFlipModel(), XYRandomRotate90Model(), N2VManipulateModel()],
+        validate_default=True,
+    )
+    """N2V compatible transforms. N2VManpulate should be the last transform."""
+
+    @field_validator("transforms")
+    @classmethod
+    def validate_transforms(
+        cls, transforms: list[N2V_TRANSFORMS_UNION]
+    ) -> list[N2V_TRANSFORMS_UNION]:
+        """
+        Validate N2VManipulate transform position in the transform list.
+
+        Parameters
+        ----------
+        transforms : list of transforms compatible with N2V
+            Transforms.
+
+        Returns
+        -------
+        list of transforms
+            Validated transforms.
+
+        Raises
+        ------
+        ValueError
+            If multiple instances of N2VManipulate are found or if it is not the last
+            transform.
+        """
+        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:
+                raise ValueError(
+                    f"{SupportedTransform.N2V_MANIPULATE} transform "
+                    f"should be the last transform."
+                )
+
+        else:
+            raise ValueError(
+                f"{SupportedTransform.N2V_MANIPULATE} transform "
+                f"is required for N2V training."
+            )
+
+        return transforms
+
+    def set_n2v2(self, use_n2v2: bool) -> None:
+        """
+        Set the N2V transform to the N2V2 version.
+
+        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_masking_strategy("median")
+        else:
+            self.set_masking_strategy("uniform")
+
+    def set_masking_strategy(self, strategy: Literal["uniform", "median"]) -> None:
+        """
+        Set masking strategy.
+
+        Parameters
+        ----------
+        strategy : "uniform" or "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 get_masking_strategy(self) -> Literal["uniform", "median"]:
+        """
+        Get N2V2 strategy.
+
+        Returns
+        -------
+        "uniform" or "median"
+            Strategy used for N2V2.
+        """
+        for transform in self.transforms:
+            if transform.name == SupportedTransform.N2V_MANIPULATE.value:
+                return transform.strategy
+
+        raise ValueError(
+            f"{SupportedTransform.N2V_MANIPULATE} transform "
+            f"is required for N2V training."
+        )
+
+    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})."
+            )
+
+    def is_using_struct_n2v(self) -> bool:
+        """
+        Check if structN2V is enabled.
+
+        Returns
+        -------
+        bool
+            Whether structN2V is enabled or not.
+        """
+        for transform in self.transforms:
+            if transform.name == SupportedTransform.N2V_MANIPULATE.value:
+                return transform.struct_mask_axis != "none"
+
+        raise ValueError(
+            f"N2V pixel manipulate transform not found in the transforms "
+            f"({self.transforms})."
+        )

careamics/config/likelihood_model.py

@@ -1,11 +1,10 @@
 """Likelihood model."""
 
-from typing import Literal, Optional, Union
+from typing import Annotated, Literal, Optional, Union
 
 import numpy as np
 import torch
 from pydantic import BaseModel, ConfigDict, PlainSerializer, PlainValidator
-from typing_extensions import Annotated
 
 from careamics.models.lvae.noise_models import (
     GaussianMixtureNoiseModel,

careamics/config/n2n_configuration.py

@@ -0,0 +1,101 @@
+"""N2N configuration."""
+
+from bioimageio.spec.generic.v0_3 import CiteEntry
+
+from careamics.config.algorithms import N2NAlgorithm
+from careamics.config.configuration import Configuration
+from careamics.config.data import DataConfig
+
+N2N = "Noise2Noise"
+
+N2N_DESCRIPTION = (
+    "Noise2Noise is a deep-learning-based algorithm that uses a U-Net "
+    "architecture to restore images. Noise2Noise is a self-supervised "
+    "algorithm that requires only noisy images to train the network. "
+    "The algorithm learns to predict the clean image from the noisy "
+    "image. Noise2Noise is particularly useful when clean images are "
+    "not available for training."
+)
+
+N2N_REF = 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",
+)
+
+
+class N2NConfiguration(Configuration):
+    """Noise2Noise configuration."""
+
+    algorithm_config: N2NAlgorithm
+    """Algorithm configuration."""
+
+    data_config: DataConfig
+    """Data configuration."""
+
+    def get_algorithm_friendly_name(self) -> str:
+        """
+        Get the algorithm friendly name.
+
+        Returns
+        -------
+        str
+            Friendly name of the algorithm.
+        """
+        return N2N
+
+    def get_algorithm_keywords(self) -> list[str]:
+        """
+        Get algorithm keywords.
+
+        Returns
+        -------
+        list[str]
+            List of keywords.
+        """
+        return [
+            "restoration",
+            "UNet",
+            "3D" if "Z" in self.data_config.axes else "2D",
+            "CAREamics",
+            "pytorch",
+            N2N,
+        ]
+
+    def get_algorithm_references(self) -> str:
+        """
+        Get the algorithm references.
+
+        This is used to generate the README of the BioImage Model Zoo export.
+
+        Returns
+        -------
+        str
+            Algorithm references.
+        """
+        return N2N_REF.text + " doi: " + N2N_REF.doi
+
+    def get_algorithm_citations(self) -> list[CiteEntry]:
+        """
+        Return a list of citation entries of the current algorithm.
+
+        This is used to generate the model description for the BioImage Model Zoo.
+
+        Returns
+        -------
+        List[CiteEntry]
+            List of citation entries.
+        """
+        return [N2N_REF]
+
+    def get_algorithm_description(self) -> str:
+        """
+        Get the algorithm description.
+
+        Returns
+        -------
+        str
+            Algorithm description.
+        """
+        return N2N_DESCRIPTION