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

careamics/__init__.py

@@ -7,7 +7,22 @@ try:
 except PackageNotFoundError:
     __version__ = "uninstalled"
 
-__all__ = ["CAREamist", "Configuration", "load_configuration", "save_configuration"]
+__all__ = [
+    "CAREamist",
+    "Configuration",
+    "algorithm_factory",
+    "configuration_factory",
+    "data_factory",
+    "load_configuration",
+    "save_configuration",
+]
 
 from .careamist import CAREamist
-from .config import Configuration, load_configuration, save_configuration
+from .config import (
+    Configuration,
+    algorithm_factory,
+    configuration_factory,
+    data_factory,
+    load_configuration,
+    save_configuration,
+)

careamics/careamist.py

@@ -13,7 +13,7 @@ from pytorch_lightning.callbacks import (
 )
 from pytorch_lightning.loggers import CSVLogger, TensorBoardLogger, WandbLogger
 
-from careamics.config import Configuration, FCNAlgorithmConfig, load_configuration
+from careamics.config import Configuration, UNetBasedAlgorithm, load_configuration
 from careamics.config.support import (
     SupportedAlgorithm,
     SupportedArchitecture,
@@ -137,7 +137,7 @@ class CAREamist:
             self.cfg = source
 
             # instantiate model
-            if isinstance(self.cfg.algorithm_config, FCNAlgorithmConfig):
+            if isinstance(self.cfg.algorithm_config, UNetBasedAlgorithm):
                 self.model = FCNModule(
                     algorithm_config=self.cfg.algorithm_config,
                 )
@@ -157,7 +157,8 @@ class CAREamist:
                 self.cfg = load_configuration(source)
 
                 # instantiate model
-                if isinstance(self.cfg.algorithm_config, FCNAlgorithmConfig):
+                # TODO call model factory here
+                if isinstance(self.cfg.algorithm_config, UNetBasedAlgorithm):
                     self.model = FCNModule(
                         algorithm_config=self.cfg.algorithm_config,
                     )  # type: ignore

careamics/cli/conf.py

@@ -3,12 +3,11 @@
 import sys
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Optional
+from typing import Annotated, Optional
 
 import click
 import typer
 import yaml
-from typing_extensions import Annotated
 
 from ..config import (
     Configuration,

careamics/cli/main.py

@@ -7,11 +7,10 @@ its implementation is contained in the conf.py file.
 """
 
 from pathlib import Path
-from typing import Optional
+from typing import Annotated, Optional
 
 import click
 import typer
-from typing_extensions import Annotated
 
 from ..careamist import CAREamist
 from . import conf

careamics/cli/utils.py

@@ -1,11 +1,11 @@
 """Utility functions for the CAREamics CLI."""
 
-from typing import Optional, Tuple
+from typing import Optional
 
 
 def handle_2D_3D_callback(
-    value: Optional[Tuple[int, int, int]]
-) -> Optional[Tuple[int, ...]]:
+    value: Optional[tuple[int, int, int]]
+) -> Optional[tuple[int, ...]]:
     """
     Callback for options that require 2D or 3D inputs.
 

careamics/config/__init__.py

@@ -1,41 +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",
-    "LVAELossConfig",
+    "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,24 +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 .loss_model import LVAELossConfig
-from .nm_model import MultiChannelNMConfig
-from .optimizer_models import LrSchedulerModel, OptimizerModel
 
-
-class VAEAlgorithmConfig(BaseModel):
-    """Algorithm configuration.
+class VAEBasedAlgorithm(BaseModel):
+    """VAE-based algorithm configuration.
 
     # TODO
 
@@ -42,7 +44,7 @@ class VAEAlgorithmConfig(BaseModel):
 
     # NOTE: these are all configs (pydantic models)
     loss: LVAELossConfig
-    model: Union[LVAEModel, CustomModel] = Field(discriminator="architecture")
+    model: LVAEModel
     noise_model: Optional[MultiChannelNMConfig] = None
     noise_model_likelihood: Optional[NMLikelihoodConfig] = None
     gaussian_likelihood: Optional[GaussianLikelihoodConfig] = None

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,12 +15,17 @@ class LVAEModel(ArchitectureModel):
     model_config = ConfigDict(validate_assignment=True, validate_default=True)
 
     architecture: Literal["LVAE"]
-    input_shape: list[int] = Field(default=(64, 64), validate_default=True)
+    """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 !!
 

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)

careamics/config/care_configuration.py

@@ -0,0 +1,100 @@
+"""CARE Pydantic configuration."""
+
+from bioimageio.spec.generic.v0_3 import CiteEntry
+
+from careamics.config.algorithms.care_algorithm_model import CAREAlgorithm
+from careamics.config.configuration import Configuration
+from careamics.config.data import DataConfig
+
+CARE = "CARE"
+
+CARE_DESCRIPTION = (
+    "Content-aware image restoration (CARE) is a deep-learning-based "
+    "algorithm that uses a U-Net architecture to restore images. CARE "
+    "is a supervised algorithm that requires pairs of noisy and "
+    "clean images to train the network. The algorithm learns to "
+    "predict the clean image from the noisy image. CARE is "
+    "particularly useful for denoising images acquired in low-light "
+    "conditions, such as fluorescence microscopy images."
+)
+CARE_REF = 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",
+)
+
+
+class CAREConfiguration(Configuration):
+    """CARE configuration."""
+
+    algorithm_config: CAREAlgorithm
+    """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 CARE
+
+    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",
+            CARE,
+        ]
+
+    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 CARE_REF.text + " doi: " + CARE_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 [CARE_REF]
+
+    def get_algorithm_description(self) -> str:
+        """
+        Get the algorithm description.
+
+        Returns
+        -------
+        str
+            Algorithm description.
+        """
+        return CARE_DESCRIPTION