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

careamics/config/__init__.py

@@ -0,0 +1,39 @@
+"""Configuration module."""
+
+__all__ = [
+    "FCNAlgorithmConfig",
+    "VAEAlgorithmConfig",
+    "DataConfig",
+    "Configuration",
+    "CheckpointModel",
+    "InferenceConfig",
+    "load_configuration",
+    "save_configuration",
+    "TrainingConfig",
+    "create_n2v_configuration",
+    "create_n2n_configuration",
+    "create_care_configuration",
+    "register_model",
+    "CustomModel",
+    "clear_custom_models",
+    "GaussianMixtureNMConfig",
+    "MultiChannelNMConfig",
+]
+from .architectures import CustomModel, clear_custom_models, register_model
+from .callback_model import CheckpointModel
+from .configuration_factory import (
+    create_care_configuration,
+    create_n2n_configuration,
+    create_n2v_configuration,
+)
+from .configuration_model import (
+    Configuration,
+    load_configuration,
+    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

@@ -0,0 +1,17 @@
+"""Deep-learning model configurations."""
+
+__all__ = [
+    "ArchitectureModel",
+    "CustomModel",
+    "UNetModel",
+    "LVAEModel",
+    "clear_custom_models",
+    "get_custom_model",
+    "register_model",
+]
+
+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

@@ -0,0 +1,37 @@
+"""Base model for the various CAREamics architectures."""
+
+from typing import Any, Dict
+
+from pydantic import BaseModel
+
+
+class ArchitectureModel(BaseModel):
+    """
+    Base Pydantic model for all model architectures.
+
+    The `model_dump` method allows removing the `architecture` key from the model.
+    """
+
+    architecture: str
+    """Name of the architecture."""
+
+    def model_dump(self, **kwargs: Any) -> Dict[str, Any]:
+        """
+        Dump the model as a dictionary, ignoring the architecture keyword.
+
+        Parameters
+        ----------
+        **kwargs : Any
+            Additional keyword arguments from Pydantic BaseModel model_dump method.
+
+        Returns
+        -------
+        dict[str, Any]
+            Model as a dictionary.
+        """
+        model_dict = super().model_dump(**kwargs)
+
+        # remove the architecture key
+        model_dict.pop("architecture")
+
+        return model_dict

careamics/config/architectures/custom_model.py

@@ -0,0 +1,162 @@
+"""Custom architecture Pydantic model."""
+
+from __future__ import annotations
+
+import inspect
+from pprint import pformat
+from typing import Any, Literal
+
+from pydantic import ConfigDict, field_validator, model_validator
+from torch.nn import Module
+from typing_extensions import Self
+
+from .architecture_model import ArchitectureModel
+from .register_model import get_custom_model
+
+
+class CustomModel(ArchitectureModel):
+    """Custom model configuration.
+
+    This Pydantic model allows storing parameters for a custom model. In order for the
+    model to be valid, the specific model needs to be registered using the
+    `register_model` decorator, and its name correctly passed to this model
+    configuration (see Examples).
+
+    Attributes
+    ----------
+    architecture : Literal["custom"]
+        Discriminator for the custom model, must be set to "custom".
+    name : str
+        Name of the custom model.
+    parameters : CustomParametersModel
+        All parameters, required for the initialization of the torch module have to be
+        passed here.
+
+    Raises
+    ------
+    ValueError
+        If the custom model `name` is unknown.
+    ValueError
+        If the custom model is not a torch Module subclass.
+    ValueError
+        If the custom model parameters are not valid.
+
+    Examples
+    --------
+    >>> from torch import nn, ones
+    >>> from careamics.config import CustomModel, register_model
+    >>> # Register a custom model
+    >>> @register_model(name="my_linear")
+    ... 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
+    ...
+    >>> # Create a configuration
+    >>> config_dict = {
+    ...     "architecture": "custom",
+    ...     "name": "my_linear",
+    ...     "in_features": 10,
+    ...     "out_features": 5,
+    ... }
+    >>> config = CustomModel(**config_dict)
+    """
+
+    # pydantic model config
+    model_config = ConfigDict(
+        extra="allow",
+    )
+
+    # discriminator used for choosing the pydantic model in Model
+    architecture: Literal["custom"]
+    """Name of the architecture."""
+
+    name: str
+    """Name of the custom model."""
+
+    @field_validator("name")
+    @classmethod
+    def custom_model_is_known(cls, value: str) -> str:
+        """Check whether the custom model is known.
+
+        Parameters
+        ----------
+        value : str
+            Name of the custom model as registered using the `@register_model`
+            decorator.
+
+        Returns
+        -------
+        str
+            The custom model name.
+        """
+        # delegate error to get_custom_model
+        model = get_custom_model(value)
+
+        # check if it is a torch Module subclass
+        if not issubclass(model, Module):
+            raise ValueError(
+                f'Retrieved class {model} with name "{value}" is not a '
+                f"torch.nn.Module subclass."
+            )
+
+        return value
+
+    @model_validator(mode="after")
+    def check_parameters(self: Self) -> Self:
+        """Validate model by instantiating the model with the parameters.
+
+        Returns
+        -------
+        Self
+            The validated model.
+        """
+        # instantiate model
+        try:
+            get_custom_model(self.name)(**self.model_dump())
+        except Exception as e:
+            raise ValueError(
+                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. Trace: "
+                f"filename: {inspect.trace()[-1].filename}, function: "
+                f"{inspect.trace()[-1].function}, line: {inspect.trace()[-1].lineno}"
+            ) from None
+
+        return self
+
+    def __str__(self) -> str:
+        """Pretty string representing the configuration.
+
+        Returns
+        -------
+        str
+            Pretty string.
+        """
+        return pformat(self.model_dump())
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        """Dump the model configuration.
+
+        Parameters
+        ----------
+        **kwargs : Any
+            Additional keyword arguments from Pydantic BaseModel model_dump method.
+
+        Returns
+        -------
+        dict[str, Any]
+            Model configuration.
+        """
+        model_dict = super().model_dump()
+
+        # remove the name key
+        model_dict.pop("name")
+
+        return model_dict

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/architectures/register_model.py

@@ -0,0 +1,103 @@
+"""Custom model registration utilities."""
+
+from typing import Callable
+
+from torch.nn import Module
+
+CUSTOM_MODELS = {}  # dictionary of custom models {"name": __class__}
+
+
+def register_model(name: str) -> Callable:
+    """Decorator used to register a torch.nn.Module class with a given `name`.
+
+    Parameters
+    ----------
+    name : str
+        Name of the model.
+
+    Returns
+    -------
+    Callable
+        Function allowing to instantiate the wrapped Module class.
+
+    Raises
+    ------
+    ValueError
+        If a model is already registered with that name.
+
+    Examples
+    --------
+    ```python
+    @register_model(name="linear")
+    class LinearModel(nn.Module):
+        def __init__(self, in_features, out_features):
+            super().__init__()
+
+            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
+    ```
+    """
+    if name is None or name == "":
+        raise ValueError("Model name cannot be empty.")
+
+    if name in CUSTOM_MODELS:
+        raise ValueError(
+            f"Model {name} already exists. Choose a different name or run "
+            f"`clear_custom_models()` to empty the registry."
+        )
+
+    def add_custom_model(model: Module) -> Module:
+        """Add a custom model to the registry and return it.
+
+        Parameters
+        ----------
+        model : Module
+            Module class to register.
+
+        Returns
+        -------
+        Module
+            The registered model.
+        """
+        # add model to the registry
+        CUSTOM_MODELS[name] = model
+
+        return model
+
+    return add_custom_model
+
+
+def get_custom_model(name: str) -> Module:
+    """Get the custom model corresponding to `name` from the registry.
+
+    Parameters
+    ----------
+    name : str
+        Name of the model to retrieve.
+
+    Returns
+    -------
+    Module
+        The requested model.
+
+    Raises
+    ------
+    ValueError
+        If the model is not registered.
+    """
+    if name not in CUSTOM_MODELS:
+        raise ValueError(
+            f"Model {name} is unknown. Have you registered it using "
+            f'@register_model("{name}") as decorator?'
+        )
+
+    return CUSTOM_MODELS[name]
+
+
+def clear_custom_models() -> None:
+    """Clear the custom models registry."""
+    # clear dictionary
+    CUSTOM_MODELS.clear()

careamics/config/architectures/unet_model.py

@@ -0,0 +1,118 @@
+"""UNet Pydantic model."""
+
+from __future__ import annotations
+
+from typing import Literal
+
+from pydantic import ConfigDict, Field, field_validator
+
+from .architecture_model import ArchitectureModel
+
+
+# TODO tests activation <-> pydantic model, test the literals!
+# TODO annotations for the json schema?
+class UNetModel(ArchitectureModel):
+    """
+    Pydantic model for a N2V(2)-compatible UNet.
+
+    Attributes
+    ----------
+    depth : int
+        Depth of the model, between 1 and 10 (default 2).
+    num_channels_init : int
+        Number of filters of the first level of the network, should be even
+        and minimum 8 (default 96).
+    """
+
+    # pydantic model config
+    model_config = ConfigDict(validate_assignment=True)
+
+    # discriminator used for choosing the pydantic model in Model
+    architecture: Literal["UNet"]
+    """Name of the architecture."""
+
+    # parameters
+    # validate_defaults allow ignoring default values in the dump if they were not set
+    conv_dims: Literal[2, 3] = Field(default=2, validate_default=True)
+    """Dimensions (2D or 3D) of the convolutional layers."""
+
+    num_classes: int = Field(default=1, ge=1, validate_default=True)
+    """Number of classes or channels in the model output."""
+
+    in_channels: int = Field(default=1, ge=1, validate_default=True)
+    """Number of channels in the input to the model."""
+
+    depth: int = Field(default=2, ge=1, le=10, validate_default=True)
+    """Number of levels in the UNet."""
+
+    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."""
+
+    final_activation: Literal[
+        "None", "Sigmoid", "Softmax", "Tanh", "ReLU", "LeakyReLU"
+    ] = Field(default="None", validate_default=True)
+    """Final activation function."""
+
+    n2v2: bool = Field(default=False, validate_default=True)
+    """Whether to use N2V2 architecture modifications, with blur pool layers and fewer
+    skip connections.
+    """
+
+    independent_channels: bool = Field(default=True, validate_default=True)
+    """Whether information is processed independently in each channel, used to train
+    channels independently."""
+
+    @field_validator("num_channels_init")
+    @classmethod
+    def validate_num_channels_init(cls, num_channels_init: int) -> int:
+        """
+        Validate that num_channels_init is even.
+
+        Parameters
+        ----------
+        num_channels_init : int
+            Number of channels.
+
+        Returns
+        -------
+        int
+            Validated number of channels.
+
+        Raises
+        ------
+        ValueError
+            If the number of channels is odd.
+        """
+        # if odd
+        if num_channels_init % 2 != 0:
+            raise ValueError(
+                f"Number of channels for the bottom layer must be even"
+                f" (got {num_channels_init})."
+            )
+
+        return num_channels_init
+
+    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.
+        """
+        if is_3D:
+            self.conv_dims = 3
+        else:
+            self.conv_dims = 2
+
+    def is_3D(self) -> bool:
+        """
+        Return whether the model is 3D or not.
+
+        Returns
+        -------
+        bool
+            Whether the model is 3D or not.
+        """
+        return self.conv_dims == 3