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

careamics/cli/utils.py

@@ -0,0 +1,29 @@
+"""Utility functions for the CAREamics CLI."""
+
+from typing import Optional
+
+
+def handle_2D_3D_callback(
+    value: Optional[tuple[int, int, int]]
+) -> Optional[tuple[int, ...]]:
+    """
+    Callback for options that require 2D or 3D inputs.
+
+    In the case of 2D, the 3rd element should be set to -1.
+
+    Parameters
+    ----------
+    value : (int, int, int)
+        Tile size value.
+
+    Returns
+    -------
+    (int, int, int) | (int, int)
+        If the last element in `value` is -1 the tuple is reduced to the first two
+        values.
+    """
+    if value is None:
+        return value
+    if value[2] == -1:
+        return value[:2]
+    return value

careamics/config/__init__.py

@@ -1,39 +1,63 @@
-"""Configuration module."""
+"""CAREamics Pydantic configuration models.
+
+To maintain clarity at the module level, we follow the following naming conventions:
+`*_model` is specific for sub-configurations (e.g. architecture, data, algorithm),
+while `*_configuration` is reserved for the main configuration models, including the
+`Configuration` base class and its algorithm-specific child classes.
+"""
 
 __all__ = [
-    "FCNAlgorithmConfig",
-    "VAEAlgorithmConfig",
-    "DataConfig",
-    "Configuration",
+    "CAREAlgorithm",
+    "CAREConfiguration",
     "CheckpointModel",
+    "Configuration",
+    "DataConfig",
+    "GaussianMixtureNMConfig",
+    "GeneralDataConfig",
     "InferenceConfig",
-    "load_configuration",
-    "save_configuration",
+    "LVAELossConfig",
+    "MultiChannelNMConfig",
+    "N2NAlgorithm",
+    "N2NConfiguration",
+    "N2VAlgorithm",
+    "N2VConfiguration",
+    "N2VDataConfig",
     "TrainingConfig",
-    "create_n2v_configuration",
-    "create_n2n_configuration",
+    "UNetBasedAlgorithm",
+    "VAEBasedAlgorithm",
+    "algorithm_factory",
+    "configuration_factory",
     "create_care_configuration",
-    "register_model",
-    "CustomModel",
-    "clear_custom_models",
-    "GaussianMixtureNMConfig",
-    "MultiChannelNMConfig",
+    "create_n2n_configuration",
+    "create_n2v_configuration",
+    "data_factory",
+    "load_configuration",
+    "save_configuration",
 ]
-from .architectures import CustomModel, clear_custom_models, register_model
+
+from .algorithms import (
+    CAREAlgorithm,
+    N2NAlgorithm,
+    N2VAlgorithm,
+    UNetBasedAlgorithm,
+    VAEBasedAlgorithm,
+)
 from .callback_model import CheckpointModel
-from .configuration_factory import (
+from .care_configuration import CAREConfiguration
+from .configuration import Configuration
+from .configuration_factories import (
+    algorithm_factory,
+    configuration_factory,
     create_care_configuration,
     create_n2n_configuration,
     create_n2v_configuration,
+    data_factory,
 )
-from .configuration_model import (
-    Configuration,
-    load_configuration,
-    save_configuration,
-)
-from .data_model import DataConfig
-from .fcn_algorithm_model import FCNAlgorithmConfig
+from .configuration_io import load_configuration, save_configuration
+from .data import DataConfig, GeneralDataConfig, N2VDataConfig
 from .inference_model import InferenceConfig
+from .loss_model import LVAELossConfig
+from .n2n_configuration import N2NConfiguration
+from .n2v_configuration import N2VConfiguration
 from .nm_model import GaussianMixtureNMConfig, MultiChannelNMConfig
 from .training_model import TrainingConfig
-from .vae_algorithm_model import VAEAlgorithmConfig

careamics/config/algorithms/__init__.py

@@ -0,0 +1,15 @@
+"""Algorithm configurations."""
+
+__all__ = [
+    "CAREAlgorithm",
+    "N2NAlgorithm",
+    "N2VAlgorithm",
+    "UNetBasedAlgorithm",
+    "VAEBasedAlgorithm",
+]
+
+from .care_algorithm_model import CAREAlgorithm
+from .n2n_algorithm_model import N2NAlgorithm
+from .n2v_algorithm_model import N2VAlgorithm
+from .unet_algorithm_model import UNetBasedAlgorithm
+from .vae_algorithm_model import VAEBasedAlgorithm

careamics/config/algorithms/care_algorithm_model.py

@@ -0,0 +1,50 @@
+"""CARE algorithm configuration."""
+
+from typing import Literal
+
+from pydantic import field_validator
+
+from careamics.config.architectures import UNetModel
+
+from .unet_algorithm_model import UNetBasedAlgorithm
+
+
+class CAREAlgorithm(UNetBasedAlgorithm):
+    """CARE algorithm configuration.
+
+    Attributes
+    ----------
+    algorithm : "care"
+        CARE Algorithm name.
+    loss : {"mae", "mse"}
+        CARE-compatible loss function.
+    """
+
+    algorithm: Literal["care"] = "care"
+    """CARE Algorithm name."""
+
+    loss: Literal["mae", "mse"] = "mae"
+    """CARE-compatible loss function."""
+
+    @classmethod
+    @field_validator("model")
+    def model_without_n2v2(cls, value: UNetModel) -> UNetModel:
+        """Validate that the model does not have the n2v2 attribute.
+
+        Parameters
+        ----------
+        value : UNetModel
+            Model to validate.
+
+        Returns
+        -------
+        UNetModel
+            The validated model.
+        """
+        if value.n2v2:
+            raise ValueError(
+                "The N2N algorithm does not support the `n2v2` attribute. "
+                "Set it to `False`."
+            )
+
+        return value

careamics/config/algorithms/n2n_algorithm_model.py

@@ -0,0 +1,42 @@
+"""N2N Algorithm configuration."""
+
+from typing import Literal
+
+from pydantic import field_validator
+
+from careamics.config.architectures import UNetModel
+
+from .unet_algorithm_model import UNetBasedAlgorithm
+
+
+class N2NAlgorithm(UNetBasedAlgorithm):
+    """N2N Algorithm configuration."""
+
+    algorithm: Literal["n2n"] = "n2n"
+    """N2N Algorithm name."""
+
+    loss: Literal["mae", "mse"] = "mae"
+    """N2N-compatible loss function."""
+
+    @classmethod
+    @field_validator("model")
+    def model_without_n2v2(cls, value: UNetModel) -> UNetModel:
+        """Validate that the model does not have the n2v2 attribute.
+
+        Parameters
+        ----------
+        value : UNetModel
+            Model to validate.
+
+        Returns
+        -------
+        UNetModel
+            The validated model.
+        """
+        if value.n2v2:
+            raise ValueError(
+                "The N2N algorithm does not support the `n2v2` attribute. "
+                "Set it to `False`."
+            )
+
+        return value

careamics/config/algorithms/n2v_algorithm_model.py

@@ -0,0 +1,35 @@
+""""N2V Algorithm configuration."""
+
+from typing import Literal
+
+from pydantic import model_validator
+from typing_extensions import Self
+
+from .unet_algorithm_model import UNetBasedAlgorithm
+
+
+class N2VAlgorithm(UNetBasedAlgorithm):
+    """N2V Algorithm configuration."""
+
+    algorithm: Literal["n2v"] = "n2v"
+    """N2V Algorithm name."""
+
+    loss: Literal["n2v"] = "n2v"
+    """N2V loss function."""
+
+    @model_validator(mode="after")
+    def algorithm_cross_validation(self: Self) -> Self:
+        """Validate the algorithm model for N2V.
+
+        Returns
+        -------
+        Self
+            The validated model.
+        """
+        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."
+            )
+
+        return self

careamics/config/algorithms/unet_algorithm_model.py

@@ -0,0 +1,88 @@
+"""UNet-based algorithm Pydantic model."""
+
+from pprint import pformat
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict
+
+from careamics.config.architectures import UNetModel
+from careamics.config.optimizer_models import LrSchedulerModel, OptimizerModel
+
+
+class UNetBasedAlgorithm(BaseModel):
+    """General UNet-based 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, and N2N algorithms. In order to train these
+    algorithms, use the corresponding configuration child classes (e.g.
+    `N2VAlgorithm`) to ensure coherent parameters (e.g. specific losses).
+
+
+    Attributes
+    ----------
+    algorithm : {"n2v", "care", "n2n"}
+        Algorithm to use.
+    loss : {"n2v", "mae", "mse"}
+        Loss function to use.
+    model : UNetModel
+        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.
+    """
+
+    # Pydantic class configuration
+    model_config = ConfigDict(
+        protected_namespaces=(),  # allows to use model_* as a field name
+        validate_assignment=True,
+        extra="allow",
+    )
+
+    # Mandatory fields
+    algorithm: Literal["n2v", "care", "n2n"]
+    """Algorithm name, as defined in SupportedAlgorithm."""
+
+    loss: Literal["n2v", "mae", "mse"]
+    """Loss function to use, as defined in SupportedLoss."""
+
+    model: UNetModel
+    """UNet model configuration."""
+
+    # Optional fields
+    optimizer: OptimizerModel = OptimizerModel()
+    """Optimizer to use, defined in SupportedOptimizer."""
+
+    lr_scheduler: LrSchedulerModel = LrSchedulerModel()
+    """Learning rate scheduler to use, defined in SupportedLrScheduler."""
+
+    def __str__(self) -> str:
+        """Pretty string representing the configuration.
+
+        Returns
+        -------
+        str
+            Pretty string.
+        """
+        return pformat(self.model_dump())
+
+    @classmethod
+    def get_compatible_algorithms(cls) -> list[str]:
+        """Get the list of compatible algorithms.
+
+        Returns
+        -------
+        list of str
+            List of compatible algorithms.
+        """
+        return ["n2v", "care", "n2n"]

careamics/config/{vae_algorithm_model.py → algorithms/vae_algorithm_model.py}

@@ -1,23 +1,26 @@
-"""Algorithm configuration."""
+"""VAE-based algorithm Pydantic model."""
 
 from __future__ import annotations
 
 from pprint import pformat
-from typing import Literal, Optional, Union
+from typing import Literal, Optional
 
-from pydantic import BaseModel, ConfigDict, Field, model_validator
+from pydantic import BaseModel, ConfigDict, model_validator
 from typing_extensions import Self
 
+from careamics.config.architectures import LVAEModel
+from careamics.config.likelihood_model import (
+    GaussianLikelihoodConfig,
+    NMLikelihoodConfig,
+)
+from careamics.config.loss_model import LVAELossConfig
+from careamics.config.nm_model import MultiChannelNMConfig
+from careamics.config.optimizer_models import LrSchedulerModel, OptimizerModel
 from careamics.config.support import SupportedAlgorithm, SupportedLoss
 
-from .architectures import CustomModel, LVAEModel
-from .likelihood_model import GaussianLikelihoodConfig, NMLikelihoodConfig
-from .nm_model import MultiChannelNMConfig
-from .optimizer_models import LrSchedulerModel, OptimizerModel
 
-
-class VAEAlgorithmConfig(BaseModel):
-    """Algorithm configuration.
+class VAEBasedAlgorithm(BaseModel):
+    """VAE-based algorithm configuration.
 
     # TODO
 
@@ -38,13 +41,13 @@ class VAEAlgorithmConfig(BaseModel):
     # TODO: Use supported Enum classes for typing?
     #   - values can still be passed as strings and they will be cast to Enum
     algorithm: Literal["musplit", "denoisplit"]
-    loss: Literal["musplit", "denoisplit", "denoisplit_musplit"]
-    model: Union[LVAEModel, CustomModel] = Field(discriminator="architecture")
 
-    # TODO: these are configs, change naming of attrs
+    # NOTE: these are all configs (pydantic models)
+    loss: LVAELossConfig
+    model: LVAEModel
     noise_model: Optional[MultiChannelNMConfig] = None
-    noise_model_likelihood_model: Optional[NMLikelihoodConfig] = None
-    gaussian_likelihood_model: Optional[GaussianLikelihoodConfig] = None
+    noise_model_likelihood: Optional[NMLikelihoodConfig] = None
+    gaussian_likelihood: Optional[GaussianLikelihoodConfig] = None
 
     # Optional fields
     optimizer: OptimizerModel = OptimizerModel()
@@ -63,13 +66,13 @@ class VAEAlgorithmConfig(BaseModel):
         """
         # musplit
         if self.algorithm == SupportedAlgorithm.MUSPLIT:
-            if self.loss != SupportedLoss.MUSPLIT:
+            if self.loss.loss_type != SupportedLoss.MUSPLIT:
                 raise ValueError(
                     f"Algorithm {self.algorithm} only supports loss `musplit`."
                 )
 
         if self.algorithm == SupportedAlgorithm.DENOISPLIT:
-            if self.loss not in [
+            if self.loss.loss_type not in [
                 SupportedLoss.DENOISPLIT,
                 SupportedLoss.DENOISPLIT_MUSPLIT,
             ]:
@@ -78,16 +81,17 @@ class VAEAlgorithmConfig(BaseModel):
                     "or `denoisplit_musplit."
                 )
             if (
-                self.loss == SupportedLoss.DENOISPLIT
+                self.loss.loss_type == SupportedLoss.DENOISPLIT
                 and self.model.predict_logvar is not None
             ):
                 raise ValueError(
                     "Algorithm `denoisplit` with loss `denoisplit` only supports "
                     "`predict_logvar` as `None`."
                 )
+
             if self.noise_model is None:
                 raise ValueError("Algorithm `denoisplit` requires a noise model.")
-        # TODO: what if algorithm is not musplit or denoisplit (HDN?)
+        # TODO: what if algorithm is not musplit or denoisplit
         return self
 
     @model_validator(mode="after")
@@ -115,14 +119,13 @@ class VAEAlgorithmConfig(BaseModel):
         Self
             The validated model.
         """
-        if self.gaussian_likelihood_model is not None:
+        if self.gaussian_likelihood is not None:
             assert (
-                self.model.predict_logvar
-                == self.gaussian_likelihood_model.predict_logvar
+                self.model.predict_logvar == self.gaussian_likelihood.predict_logvar
             ), (
                 f"Model `predict_logvar` ({self.model.predict_logvar}) must match "
                 "Gaussian likelihood model `predict_logvar` "
-                f"({self.gaussian_likelihood_model.predict_logvar}).",
+                f"({self.gaussian_likelihood.predict_logvar}).",
             )
         return self
 

careamics/config/architectures/__init__.py

@@ -1,17 +1,7 @@
 """Deep-learning model configurations."""
 
-__all__ = [
-    "ArchitectureModel",
-    "CustomModel",
-    "UNetModel",
-    "LVAEModel",
-    "clear_custom_models",
-    "get_custom_model",
-    "register_model",
-]
+__all__ = ["ArchitectureModel", "LVAEModel", "UNetModel"]
 
 from .architecture_model import ArchitectureModel
-from .custom_model import CustomModel
 from .lvae_model import LVAEModel
-from .register_model import clear_custom_models, get_custom_model, register_model
 from .unet_model import UNetModel

careamics/config/architectures/architecture_model.py

@@ -1,6 +1,6 @@
 """Base model for the various CAREamics architectures."""
 
-from typing import Any, Dict
+from typing import Any
 
 from pydantic import BaseModel
 
@@ -15,7 +15,7 @@ class ArchitectureModel(BaseModel):
     architecture: str
     """Name of the architecture."""
 
-    def model_dump(self, **kwargs: Any) -> Dict[str, Any]:
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
         """
         Dump the model as a dictionary, ignoring the architecture keyword.
 
@@ -26,7 +26,7 @@ class ArchitectureModel(BaseModel):
 
         Returns
         -------
-        dict[str, Any]
+        {str: Any}
             Model as a dictionary.
         """
         model_dict = super().model_dump(**kwargs)

careamics/config/architectures/lvae_model.py

@@ -15,9 +15,21 @@ class LVAEModel(ArchitectureModel):
     model_config = ConfigDict(validate_assignment=True, validate_default=True)
 
     architecture: Literal["LVAE"]
-    input_shape: int = Field(default=64, ge=8, le=1024)
-    multiscale_count: int = Field(default=5)  # TODO clarify
-    # 0 - off, len(z_dims) + 1 # TODO can/should be le to z_dims len + 1
+    """Name of the architecture."""
+
+    input_shape: list[int] = Field(default=[64, 64], validate_default=True)
+    """Shape of the input patch (C, Z, Y, X) or (C, Y, X) if the data is 2D."""
+
+    encoder_conv_strides: list = Field(default=[2, 2], validate_default=True)
+
+    # TODO make this per hierarchy step ?
+    decoder_conv_strides: list = Field(default=[2, 2], validate_default=True)
+    """Dimensions (2D or 3D) of the convolutional layers."""
+
+    multiscale_count: int = Field(default=1)
+    # TODO there should be a check for multiscale_count in dataset !!
+
+    # 1 - off, len(z_dims) + 1 # TODO Consider starting from 0
     z_dims: list = Field(default=[128, 128, 128, 128])
     output_channels: int = Field(default=1, ge=1)
     encoder_n_filters: int = Field(default=64, ge=8, le=1024)
@@ -31,10 +43,90 @@ class LVAEModel(ArchitectureModel):
     )
 
     predict_logvar: Literal[None, "pixelwise"] = None
+    analytical_kl: bool = Field(default=False)
 
-    analytical_kl: bool = Field(
-        default=False,
-    )
+    @model_validator(mode="after")
+    def validate_conv_strides(self: Self) -> Self:
+        """
+        Validate the convolutional strides.
+
+        Returns
+        -------
+        list
+            Validated strides.
+
+        Raises
+        ------
+        ValueError
+            If the number of strides is not 2.
+        """
+        if len(self.encoder_conv_strides) < 2 or len(self.encoder_conv_strides) > 3:
+            raise ValueError(
+                f"Strides must be 2 or 3 (got {len(self.encoder_conv_strides)})."
+            )
+
+        if len(self.decoder_conv_strides) < 2 or len(self.decoder_conv_strides) > 3:
+            raise ValueError(
+                f"Strides must be 2 or 3 (got {len(self.decoder_conv_strides)})."
+            )
+
+        # adding 1 to encoder strides for the number of input channels
+        if len(self.input_shape) != len(self.encoder_conv_strides):
+            raise ValueError(
+                f"Input dimensions must be equal to the number of encoder conv strides"
+                f" (got {len(self.input_shape)} and {len(self.encoder_conv_strides)})."
+            )
+
+        if len(self.encoder_conv_strides) < len(self.decoder_conv_strides):
+            raise ValueError(
+                f"Decoder can't be 3D when encoder is 2D (got"
+                f" {len(self.encoder_conv_strides)} and"
+                f"{len(self.decoder_conv_strides)})."
+            )
+
+        if any(s < 1 for s in self.encoder_conv_strides) or any(
+            s < 1 for s in self.decoder_conv_strides
+        ):
+            raise ValueError(
+                f"All strides must be greater or equal to 1"
+                f"(got {self.encoder_conv_strides} and {self.decoder_conv_strides})."
+            )
+        # TODO: validate max stride size ?
+        return self
+
+    @field_validator("input_shape")
+    @classmethod
+    def validate_input_shape(cls, input_shape: list) -> list:
+        """
+        Validate the input shape.
+
+        Parameters
+        ----------
+        input_shape : list
+            Shape of the input patch.
+
+        Returns
+        -------
+        list
+            Validated input shape.
+
+        Raises
+        ------
+        ValueError
+            If the number of dimensions is not 3 or 4.
+        """
+        if len(input_shape) < 2 or len(input_shape) > 3:
+            raise ValueError(
+                f"Number of input dimensions must be 2 for 2D data 3 for 3D"
+                f"(got {len(input_shape)})."
+            )
+
+        if any(s < 1 for s in input_shape):
+            raise ValueError(
+                f"Input shape must be greater than 1 in all dimensions"
+                f"(got {input_shape})."
+            )
+        return input_shape
 
     @field_validator("encoder_n_filters")
     @classmethod
@@ -124,27 +216,20 @@ class LVAEModel(ArchitectureModel):
         return z_dims
 
     @model_validator(mode="after")
-    def validate_multiscale_count(cls, self: Self) -> Self:
+    def validate_multiscale_count(self: Self) -> Self:
         """
         Validate the multiscale count.
 
-        Parameters
-        ----------
-        self : Self
-            The model.
-
         Returns
         -------
         Self
             The validated model.
         """
-        # if self.multiscale_count != 0:
-        #     if self.multiscale_count != len(self.z_dims) - 1:
-        #         raise ValueError(
-        #             f"Multiscale count must be 0 or equal to the number of Z "
-        #             f"dims - 1 (got {self.multiscale_count} and {len(self.z_dims)})."
-        #         )
-
+        if self.multiscale_count < 1 or self.multiscale_count > len(self.z_dims) + 1:
+            raise ValueError(
+                f"Multiscale count must be 1 for LC off or less or equal to the number"
+                f" of Z dims + 1 (got {self.multiscale_count} and {len(self.z_dims)})."
+            )
         return self
 
     def set_3D(self, is_3D: bool) -> None:
@@ -156,7 +241,10 @@ class LVAEModel(ArchitectureModel):
         is_3D : bool
             Whether the algorithm is 3D or not.
         """
-        raise NotImplementedError("VAE is not implemented yet.")
+        if is_3D:
+            self.conv_dims = 3
+        else:
+            self.conv_dims = 2
 
     def is_3D(self) -> bool:
         """
@@ -167,4 +255,4 @@ class LVAEModel(ArchitectureModel):
         bool
             Whether the model is 3D or not.
         """
-        raise NotImplementedError("VAE is not implemented yet.")
+        return self.conv_dims == 3

careamics/config/architectures/unet_model.py

@@ -48,6 +48,7 @@ class UNetModel(ArchitectureModel):
     num_channels_init: int = Field(default=32, ge=8, le=1024, validate_default=True)
     """Number of convolutional filters in the first layer of the UNet."""
 
+    # TODO we are not using this, so why make it a choice?
     final_activation: Literal[
         "None", "Sigmoid", "Softmax", "Tanh", "ReLU", "LeakyReLU"
     ] = Field(default="None", validate_default=True)