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

careamics/careamist.py

@@ -13,10 +13,7 @@ from pytorch_lightning.callbacks import (
 )
 from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger
 
-from careamics.config import (
-    Configuration,
-    load_configuration,
-)
+from careamics.config import Configuration, FCNAlgorithmConfig, load_configuration
 from careamics.config.support import (
     SupportedAlgorithm,
     SupportedArchitecture,
@@ -25,7 +22,7 @@ from careamics.config.support import (
 )
 from careamics.dataset.dataset_utils import reshape_array
 from careamics.lightning import (
-    CAREamicsModule,
+    FCNModule,
     HyperParametersCallback,
     PredictDataModule,
     ProgressBarCallback,
@@ -148,9 +145,12 @@ class CAREamist:
             self.cfg = source
 
             # instantiate model
-            self.model = CAREamicsModule(
-                algorithm_config=self.cfg.algorithm_config,
-            )
+            if isinstance(self.cfg.algorithm_config, FCNAlgorithmConfig):
+                self.model = FCNModule(
+                    algorithm_config=self.cfg.algorithm_config,
+                )
+            else:
+                raise NotImplementedError("Architecture not supported.")
 
         # path to configuration file or model
         else:
@@ -164,9 +164,12 @@ class CAREamist:
                 self.cfg = load_configuration(source)
 
                 # instantiate model
-                self.model = CAREamicsModule(
-                    algorithm_config=self.cfg.algorithm_config,
-                )
+                if isinstance(self.cfg.algorithm_config, FCNAlgorithmConfig):
+                    self.model = FCNModule(
+                        algorithm_config=self.cfg.algorithm_config,
+                    )  # type: ignore
+                else:
+                    raise NotImplementedError("Architecture not supported.")
 
             # attempt loading a pre-trained model
             else:

careamics/config/__init__.py

@@ -1,7 +1,8 @@
 """Configuration module."""
 
 __all__ = [
-    "AlgorithmConfig",
+    "FCNAlgorithmConfig",
+    "VAEAlgorithmConfig",
     "DataConfig",
     "Configuration",
     "CheckpointModel",
@@ -15,9 +16,9 @@ __all__ = [
     "register_model",
     "CustomModel",
     "clear_custom_models",
+    "GaussianMixtureNMConfig",
+    "MultiChannelNMConfig",
 ]
-
-from .algorithm_model import AlgorithmConfig
 from .architectures import CustomModel, clear_custom_models, register_model
 from .callback_model import CheckpointModel
 from .configuration_factory import (
@@ -31,5 +32,8 @@ from .configuration_model import (
     save_configuration,
 )
 from .data_model import DataConfig
+from .fcn_algorithm_model import FCNAlgorithmConfig
 from .inference_model import InferenceConfig
+from .nm_model import GaussianMixtureNMConfig, MultiChannelNMConfig
 from .training_model import TrainingConfig
+from .vae_algorithm_model import VAEAlgorithmConfig

careamics/config/architectures/__init__.py

@@ -4,7 +4,7 @@ __all__ = [
     "ArchitectureModel",
     "CustomModel",
     "UNetModel",
-    "VAEModel",
+    "LVAEModel",
     "clear_custom_models",
     "get_custom_model",
     "register_model",
@@ -12,6 +12,6 @@ __all__ = [
 
 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
-from .vae_model import VAEModel

careamics/config/architectures/architecture_model.py

@@ -27,7 +27,7 @@ class ArchitectureModel(BaseModel):
         Returns
         -------
         dict[str, Any]
-            Model as a dictionnary.
+            Model as a dictionary.
         """
         model_dict = super().model_dump(**kwargs)
 

careamics/config/architectures/custom_model.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import inspect
 from pprint import pformat
 from typing import Any, Literal
 
@@ -23,12 +24,13 @@ class CustomModel(ArchitectureModel):
 
     Attributes
     ----------
-    architecture : Literal["Custom"]
-        Discriminator for the custom model, must be set to "Custom".
+    architecture : Literal["custom"]
+        Discriminator for the custom model, must be set to "custom".
     name : str
         Name of the custom model.
     parameters : CustomParametersModel
-        Parameters of the custom model.
+        All parameters, required for the initialization of the torch module have to be
+        passed here.
 
     Raises
     ------
@@ -57,7 +59,7 @@ class CustomModel(ArchitectureModel):
     ...
     >>> # Create a configuration
     >>> config_dict = {
-    ...     "architecture": "Custom",
+    ...     "architecture": "custom",
     ...     "name": "my_linear",
     ...     "in_features": 10,
     ...     "out_features": 5,
@@ -71,10 +73,9 @@ class CustomModel(ArchitectureModel):
     )
 
     # discriminator used for choosing the pydantic model in Model
-    architecture: Literal["Custom"]
+    architecture: Literal["custom"]
     """Name of the architecture."""
 
-    # name of the custom model
     name: str
     """Name of the custom model."""
 
@@ -120,10 +121,12 @@ class CustomModel(ArchitectureModel):
             get_custom_model(self.name)(**self.model_dump())
         except Exception as e:
             raise ValueError(
-                f"error while passing parameters to the model {e}. Verify that all "
+                f"while passing parameters to the model {e}. Verify that all "
                 f"mandatory parameters are provided, and that either the {e} accepts "
                 f"*args and **kwargs in its __init__() method, or that no additional"
-                f"parameter is provided."
+                f"parameter is provided. Trace: "
+                f"filename: {inspect.trace()[-1].filename}, function: "
+                f"{inspect.trace()[-1].function}, line: {inspect.trace()[-1].lineno}"
             ) from None
 
         return self

careamics/config/architectures/lvae_model.py

@@ -0,0 +1,174 @@
+"""LVAE Pydantic model."""
+
+from typing import Literal
+
+from pydantic import ConfigDict, Field, field_validator, model_validator
+from typing_extensions import Self
+
+from .architecture_model import ArchitectureModel
+
+
+# TODO: it is quite confusing to call this LVAEModel, as it is basically a config
+class LVAEModel(ArchitectureModel):
+    """LVAE model."""
+
+    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
+    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)
+    decoder_n_filters: int = Field(default=64, ge=8, le=1024)
+    encoder_dropout: float = Field(default=0.1, ge=0.0, le=0.9)
+    decoder_dropout: float = Field(default=0.1, ge=0.0, le=0.9)
+    nonlinearity: Literal[
+        "None", "Sigmoid", "Softmax", "Tanh", "ReLU", "LeakyReLU", "ELU"
+    ] = Field(
+        default="ELU",
+    )
+
+    predict_logvar: Literal[None, "pixelwise"] = None
+
+    # TODO this parameter is exessive -> Remove & refactor
+    enable_noise_model: bool = Field(
+        default=True,
+    )
+    analytical_kl: bool = Field(
+        default=False,
+    )
+
+    @field_validator("encoder_n_filters")
+    @classmethod
+    def validate_encoder_even(cls, encoder_n_filters: int) -> int:
+        """
+        Validate that num_channels_init is even.
+
+        Parameters
+        ----------
+        encoder_n_filters : int
+            Number of channels.
+
+        Returns
+        -------
+        int
+            Validated number of channels.
+
+        Raises
+        ------
+        ValueError
+            If the number of channels is odd.
+        """
+        # if odd
+        if encoder_n_filters % 2 != 0:
+            raise ValueError(
+                f"Number of channels for the bottom layer must be even"
+                f" (got {encoder_n_filters})."
+            )
+
+        return encoder_n_filters
+
+    @field_validator("decoder_n_filters")
+    @classmethod
+    def validate_decoder_even(cls, decoder_n_filters: int) -> int:
+        """
+        Validate that num_channels_init is even.
+
+        Parameters
+        ----------
+        decoder_n_filters : int
+            Number of channels.
+
+        Returns
+        -------
+        int
+            Validated number of channels.
+
+        Raises
+        ------
+        ValueError
+            If the number of channels is odd.
+        """
+        # if odd
+        if decoder_n_filters % 2 != 0:
+            raise ValueError(
+                f"Number of channels for the bottom layer must be even"
+                f" (got {decoder_n_filters})."
+            )
+
+        return decoder_n_filters
+
+    @field_validator("z_dims")
+    def validate_z_dims(cls, z_dims: tuple) -> tuple:
+        """
+        Validate the z_dims.
+
+        Parameters
+        ----------
+        z_dims : tuple
+            Tuple of z dimensions.
+
+        Returns
+        -------
+        tuple
+            Validated z dimensions.
+
+        Raises
+        ------
+        ValueError
+            If the number of z dimensions is not 4.
+        """
+        if len(z_dims) < 2:
+            raise ValueError(
+                f"Number of z dimensions must be at least 2 (got {len(z_dims)})."
+            )
+
+        return z_dims
+
+    @model_validator(mode="after")
+    def validate_multiscale_count(cls, 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)})."
+        #         )
+
+        return self
+
+    def set_3D(self, is_3D: bool) -> None:
+        """
+        Set 3D model by setting the `conv_dims` parameters.
+
+        Parameters
+        ----------
+        is_3D : bool
+            Whether the algorithm is 3D or not.
+        """
+        raise NotImplementedError("VAE is not implemented yet.")
+
+    def is_3D(self) -> bool:
+        """
+        Return whether the model is 3D or not.
+
+        Returns
+        -------
+        bool
+            Whether the model is 3D or not.
+        """
+        raise NotImplementedError("VAE is not implemented yet.")

careamics/config/configuration_factory.py

@@ -2,10 +2,10 @@
 
 from typing import Any, Dict, List, Literal, Optional
 
-from .algorithm_model import AlgorithmConfig
 from .architectures import UNetModel
 from .configuration_model import Configuration
 from .data_model import DataConfig
+from .fcn_algorithm_model import FCNAlgorithmConfig
 from .support import (
     SupportedAlgorithm,
     SupportedArchitecture,
@@ -16,7 +16,9 @@ from .support import (
 from .training_model import TrainingConfig
 
 
+# TODO rename ?
 def _create_supervised_configuration(
+    algorithm_type: Literal["fcn"],
     algorithm: Literal["care", "n2n"],
     experiment_name: str,
     data_type: Literal["array", "tiff", "custom"],
@@ -37,6 +39,8 @@ def _create_supervised_configuration(
 
     Parameters
     ----------
+    algorithm_type : Literal["fcn"]
+        Type of the algorithm.
     algorithm : Literal["care", "n2n"]
         Algorithm to use.
     experiment_name : str
@@ -97,7 +101,8 @@ def _create_supervised_configuration(
     )
 
     # algorithm model
-    algorithm = AlgorithmConfig(
+    algorithm = FCNAlgorithmConfig(
+        algorithm_type=algorithm_type,
         algorithm=algorithm,
         loss=loss,
         model=unet_model,
@@ -215,6 +220,7 @@ def create_care_configuration(
         n_channels_out = n_channels_in
 
     return _create_supervised_configuration(
+        algorithm_type="fcn",
         algorithm="care",
         experiment_name=experiment_name,
         data_type=data_type,
@@ -304,6 +310,7 @@ def create_n2n_configuration(
         n_channels_out = n_channels_in
 
     return _create_supervised_configuration(
+        algorithm_type="fcn",
         algorithm="n2n",
         experiment_name=experiment_name,
         data_type=data_type,
@@ -514,7 +521,8 @@ def create_n2v_configuration(
     )
 
     # algorithm model
-    algorithm = AlgorithmConfig(
+    algorithm = FCNAlgorithmConfig(
+        algorithm_type="fcn",
         algorithm=SupportedAlgorithm.N2V.value,
         loss=SupportedLoss.N2V.value,
         model=unet_model,

careamics/config/configuration_model.py

@@ -9,11 +9,11 @@ from typing import Literal, Union
 
 import yaml
 from bioimageio.spec.generic.v0_3 import CiteEntry
-from pydantic import BaseModel, ConfigDict, field_validator, model_validator
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 from typing_extensions import Self
 
-from .algorithm_model import AlgorithmConfig
 from .data_model import DataConfig
+from .fcn_algorithm_model import FCNAlgorithmConfig
 from .references import (
     CARE,
     CUSTOM,
@@ -39,6 +39,7 @@ from .training_model import TrainingConfig
 from .transformations.n2v_manipulate_model import (
     N2VManipulateModel,
 )
+from .vae_algorithm_model import VAEAlgorithmConfig
 
 
 class Configuration(BaseModel):
@@ -123,6 +124,7 @@ class Configuration(BaseModel):
     >>> config_dict = {
     ...         "experiment_name": "N2V_experiment",
     ...         "algorithm_config": {
+    ...             "algorithm_type": "fcn",
     ...             "algorithm": "n2v",
     ...             "loss": "n2v",
     ...             "model": {
@@ -155,7 +157,9 @@ class Configuration(BaseModel):
     """Name of the experiment, used to name logs and checkpoints."""
 
     # Sub-configurations
-    algorithm_config: AlgorithmConfig
+    algorithm_config: Union[FCNAlgorithmConfig, VAEAlgorithmConfig] = Field(
+        discriminator="algorithm_type"
+    )
     """Algorithm configuration, holding all parameters required to configure the
     model."""
 

careamics/config/data_model.py

@@ -5,12 +5,14 @@ from __future__ import annotations
 from pprint import pformat
 from typing import Any, Literal, Optional, Union
 
+import numpy as np
 from numpy.typing import NDArray
 from pydantic import (
     BaseModel,
     ConfigDict,
     Discriminator,
     Field,
+    PlainSerializer,
     field_validator,
     model_validator,
 )
@@ -22,6 +24,30 @@ from .transformations.xy_flip_model import XYFlipModel
 from .transformations.xy_random_rotate90_model import XYRandomRotate90Model
 from .validators import check_axes_validity, patch_size_ge_than_8_power_of_2
 
+
+def np_float_to_scientific_str(x: float) -> str:
+    """Return a string scientific representation of a float.
+
+    In particular, this method is used to serialize floats to strings, allowing
+    numpy.float32 to be passed in the Pydantic model and written to a yaml file as str.
+
+    Parameters
+    ----------
+    x : float
+        Input value.
+
+    Returns
+    -------
+    str
+        Scientific string representation of the input value.
+    """
+    return np.format_float_scientific(x, precision=7)
+
+
+Float = Annotated[float, PlainSerializer(np_float_to_scientific_str, return_type=str)]
+"""Annotated float type, used to serialize floats to strings."""
+
+
 TRANSFORMS_UNION = Annotated[
     Union[
         XYFlipModel,
@@ -30,6 +56,7 @@ TRANSFORMS_UNION = Annotated[
     ],
     Discriminator("name"),  # used to tell the different transform models apart
 ]
+"""Available transforms in CAREamics."""
 
 
 class DataConfig(BaseModel):
@@ -94,20 +121,20 @@ class DataConfig(BaseModel):
     """Batch size for training."""
 
     # Optional fields
-    image_means: Optional[list[float]] = Field(
+    image_means: Optional[list[Float]] = Field(
         default=None, min_length=0, max_length=32
     )
     """Means of the data across channels, used for normalization."""
 
-    image_stds: Optional[list[float]] = Field(default=None, min_length=0, max_length=32)
+    image_stds: Optional[list[Float]] = Field(default=None, min_length=0, max_length=32)
     """Standard deviations of the data across channels, used for normalization."""
 
-    target_means: Optional[list[float]] = Field(
+    target_means: Optional[list[Float]] = Field(
         default=None, min_length=0, max_length=32
     )
     """Means of the target data across channels, used for normalization."""
 
-    target_stds: Optional[list[float]] = Field(
+    target_stds: Optional[list[Float]] = Field(
         default=None, min_length=0, max_length=32
     )
     """Standard deviations of the target data across channels, used for
@@ -265,9 +292,7 @@ class DataConfig(BaseModel):
         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."
-            )
+            raise ValueError("Mean and std must be specified for each input channel.")
 
         if (self.target_means and not self.target_stds) or (
             self.target_stds and not self.target_means
@@ -380,7 +405,7 @@ class DataConfig(BaseModel):
 
         Parameters
         ----------
-        image_means : numpy.ndarray ,tuple or list
+        image_means : numpy.ndarray, tuple or list
             Mean values for normalization.
         image_stds : numpy.ndarray, tuple or list
             Standard deviation values for normalization.

careamics/config/{algorithm_model.py → fcn_algorithm_model.py}

@@ -1,6 +1,4 @@
-"""Algorithm configuration."""
-
-from __future__ import annotations
+"""Module containing `FCNAlgorithmConfig` class."""
 
 from pprint import pformat
 from typing import Literal, Union
@@ -8,11 +6,11 @@ from typing import Literal, Union
 from pydantic import BaseModel, ConfigDict, Field, model_validator
 from typing_extensions import Self
 
-from .architectures import CustomModel, UNetModel, VAEModel
-from .optimizer_models import LrSchedulerModel, OptimizerModel
+from careamics.config.architectures import CustomModel, UNetModel
+from careamics.config.optimizer_models import LrSchedulerModel, OptimizerModel
 
 
-class AlgorithmConfig(BaseModel):
+class FCNAlgorithmConfig(BaseModel):
     """Algorithm configuration.
 
     This Pydantic model validates the parameters governing the components of the
@@ -30,7 +28,7 @@ class AlgorithmConfig(BaseModel):
         Algorithm to use.
     loss : Literal["n2v", "mae", "mse"]
         Loss function to use.
-    model : Union[UNetModel, VAEModel, CustomModel]
+    model : Union[UNetModel, LVAEModel, CustomModel]
         Model architecture to use.
     optimizer : OptimizerModel, optional
         Optimizer to use.
@@ -47,66 +45,51 @@ class AlgorithmConfig(BaseModel):
     Examples
     --------
     Minimum example:
-    >>> from careamics.config import AlgorithmConfig
+    >>> from careamics.config import FCNAlgorithmConfig
     >>> config_dict = {
     ...     "algorithm": "n2v",
+    ...     "algorithm_type": "fcn",
     ...     "loss": "n2v",
     ...     "model": {
     ...         "architecture": "UNet",
     ...     }
     ... }
-    >>> config = AlgorithmConfig(**config_dict)
-
-    Using a custom model:
-    >>> from torch import nn, ones
-    >>> from careamics.config import AlgorithmConfig, register_model
-    ...
-    >>> @register_model(name="linear_model")
-    ... class LinearModel(nn.Module):
-    ...    def __init__(self, in_features, out_features, *args, **kwargs):
-    ...        super().__init__()
-    ...        self.in_features = in_features
-    ...        self.out_features = out_features
-    ...        self.weight = nn.Parameter(ones(in_features, out_features))
-    ...        self.bias = nn.Parameter(ones(out_features))
-    ...    def forward(self, input):
-    ...        return (input @ self.weight) + self.bias
-    ...
-    >>> config_dict = {
-    ...     "algorithm": "custom",
-    ...     "loss": "mse",
-    ...     "model": {
-    ...         "architecture": "Custom",
-    ...         "name": "linear_model",
-    ...         "in_features": 10,
-    ...         "out_features": 5,
-    ...     }
-    ... }
-    >>> config = AlgorithmConfig(**config_dict)
+    >>> config = FCNAlgorithmConfig(**config_dict)
     """
 
     # 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", "custom"]  # defined in SupportedAlgorithm
-    """Name of the algorithm, as defined in SupportedAlgorithm."""
+    # defined in SupportedAlgorithm
+    algorithm_type: Literal["fcn"]
+    """Algorithm type must be `fcn` (fully convolutional network) to differentiate this
+    configuration from LVAE."""
+
+    algorithm: Literal["n2v", "care", "n2n", "custom"]
+    """Name of the algorithm, as defined in SupportedAlgorithm. Use `custom` for custom
+    model architecture."""
 
     loss: Literal["n2v", "mae", "mse"]
     """Loss function to use, as defined in SupportedLoss."""
 
-    model: Union[UNetModel, VAEModel, CustomModel] = Field(discriminator="architecture")
-    """Model architecture to use, defined in SupportedArchitecture."""
+    model: Union[UNetModel, CustomModel] = Field(discriminator="architecture")
+    """Model architecture to use, along with its parameters. Compatible architectures
+    are defined in SupportedArchitecture, and their Pydantic models in
+    `careamics.config.architectures`."""
+    # TODO supported architectures are now all the architectures but does not warn users
+    # of the compatibility with the algorithm
 
     # Optional fields
     optimizer: OptimizerModel = OptimizerModel()
     """Optimizer to use, defined in SupportedOptimizer."""
 
     lr_scheduler: LrSchedulerModel = LrSchedulerModel()
-    """Learning rate scheduler to use, defined in SupportedScheduler."""
+    """Learning rate scheduler to use, defined in SupportedLrScheduler."""
 
     @model_validator(mode="after")
     def algorithm_cross_validation(self: Self) -> Self:
@@ -146,8 +129,10 @@ class AlgorithmConfig(BaseModel):
             if self.loss == "n2v":
                 raise ValueError("Supervised algorithms do not support loss `n2v`.")
 
-        if isinstance(self.model, VAEModel):
-            raise ValueError("VAE are currently not implemented.")
+        if (self.algorithm == "custom") != (self.model.architecture == "custom"):
+            raise ValueError(
+                "Algorithm and model architecture must be both `custom` or not."
+            )
 
         return self