careamics 0.0.4.2__py3-none-any.whl → 0.0.6__py3-none-any.whl

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

@@ -2,8 +2,9 @@
 
 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
 
 import numpy as np
 from numpy.typing import NDArray
@@ -15,11 +16,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 +45,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,22 +87,18 @@ 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."""
 
     dataloader_params: Optional[dict] = None
     """Dictionary of PyTorch dataloader parameters."""
@@ -210,48 +167,6 @@ class DataConfig(BaseModel):
 
         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:
         """
@@ -350,32 +265,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 +319,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, Field, PlainSerializer, PlainValidator
-from typing_extensions import Annotated
+from pydantic import BaseModel, ConfigDict, PlainSerializer, PlainValidator
 
 from careamics.models.lvae.noise_models import (
     GaussianMixtureNoiseModel,
@@ -41,7 +40,12 @@ class GaussianLikelihoodConfig(BaseModel):
 
 
 class NMLikelihoodConfig(BaseModel):
-    """Noise model likelihood configuration."""
+    """Noise model likelihood configuration.
+
+    NOTE: we need to define the data mean and std here because the noise model
+    is trained on not-normalized data. Hence, we need to unnormalize the model
+    output to compute the noise model likelihood.
+    """
 
     model_config = ConfigDict(validate_assignment=True, arbitrary_types_allowed=True)
 
@@ -54,7 +58,3 @@ class NMLikelihoodConfig(BaseModel):
     data_std: Tensor = torch.ones(1)
     """The standard deviation of the data, used to unnormalize data for noise
     model evaluation. Shape is (target_ch,) (or (1, target_ch, [1], 1, 1))."""
-
-    # TODO: serialization/deserialization for this
-    noise_model: Optional[NoiseModel] = Field(default=None, exclude=True)
-    """The noise model instance used to compute the likelihood."""

careamics/config/loss_model.py

@@ -0,0 +1,56 @@
+"""Configuration classes for LVAE losses."""
+
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict
+
+
+class KLLossConfig(BaseModel):
+    """KL loss configuration."""
+
+    model_config = ConfigDict(validate_assignment=True, validate_default=True)
+
+    loss_type: Literal["kl", "kl_restricted"] = "kl"
+    """Type of KL divergence used as KL loss."""
+    rescaling: Literal["latent_dim", "image_dim"] = "latent_dim"
+    """Rescaling of the KL loss."""
+    aggregation: Literal["sum", "mean"] = "mean"
+    """Aggregation of the KL loss across different layers."""
+    free_bits_coeff: float = 0.0
+    """Free bits coefficient for the KL loss."""
+    annealing: bool = False
+    """Whether to apply KL loss annealing."""
+    start: int = -1
+    """Epoch at which KL loss annealing starts."""
+    annealtime: int = 10
+    """Number of epochs for which KL loss annealing is applied."""
+    current_epoch: int = 0
+    """Current epoch in the training loop."""
+
+
+class LVAELossConfig(BaseModel):
+    """LVAE loss configuration."""
+
+    model_config = ConfigDict(
+        validate_assignment=True, validate_default=True, arbitrary_types_allowed=True
+    )
+
+    loss_type: Literal["musplit", "denoisplit", "denoisplit_musplit"]
+    """Type of loss to use for LVAE."""
+
+    reconstruction_weight: float = 1.0
+    """Weight for the reconstruction loss in the total net loss
+    (i.e., `net_loss = reconstruction_weight * rec_loss + kl_weight * kl_loss`)."""
+    kl_weight: float = 1.0
+    """Weight for the KL loss in the total net loss.
+    (i.e., `net_loss = reconstruction_weight * rec_loss + kl_weight * kl_loss`)."""
+    musplit_weight: float = 0.1
+    """Weight for the muSplit loss (used in the muSplit-denoiSplit loss)."""
+    denoisplit_weight: float = 0.9
+    """Weight for the denoiSplit loss (used in the muSplit-deonoiSplit loss)."""
+    kl_params: KLLossConfig = KLLossConfig()
+    """KL loss configuration."""
+
+    # TODO: remove?
+    non_stochastic: bool = False
+    """Whether to sample latents and compute KL."""