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

careamics/config/configuration.py

@@ -0,0 +1,354 @@
+"""Pydantic CAREamics configuration."""
+
+from __future__ import annotations
+
+import re
+from pprint import pformat
+from typing import Any, Literal, Union
+
+from bioimageio.spec.generic.v0_3 import CiteEntry
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
+from typing_extensions import Self
+
+from careamics.config.algorithms import UNetBasedAlgorithm, VAEBasedAlgorithm
+from careamics.config.data import GeneralDataConfig
+from careamics.config.training_model import TrainingConfig
+
+
+class Configuration(BaseModel):
+    """
+    CAREamics configuration.
+
+    The configuration defines all parameters used to build and train a CAREamics model.
+    These parameters are validated to ensure that they are compatible with each other.
+
+    It contains three sub-configurations:
+
+    - AlgorithmModel: configuration for the algorithm training, which includes the
+        architecture, loss function, optimizer, and other hyperparameters.
+    - DataModel: configuration for the dataloader, which includes the type of data,
+        transformations, mean/std and other parameters.
+    - TrainingModel: configuration for the training, which includes the number of
+        epochs or the callbacks.
+
+    Attributes
+    ----------
+    experiment_name : str
+        Name of the experiment, used when saving logs and checkpoints.
+    algorithm : AlgorithmModel
+        Algorithm configuration.
+    data : DataModel
+        Data configuration.
+    training : TrainingModel
+        Training configuration.
+
+    Methods
+    -------
+    set_3D(is_3D: bool, axes: str, patch_size: List[int]) -> None
+        Switch configuration between 2D and 3D.
+    model_dump(
+        exclude_defaults: bool = False, exclude_none: bool = True, **kwargs: Dict
+        ) -> Dict
+        Export configuration to a dictionary.
+
+    Raises
+    ------
+    ValueError
+        Configuration parameter type validation errors.
+    ValueError
+        If the experiment name contains invalid characters or is empty.
+    ValueError
+        If the algorithm is 3D but there is not "Z" in the data axes, or 2D algorithm
+        with "Z" in data axes.
+    ValueError
+        Algorithm, data or training validation errors.
+
+    Notes
+    -----
+    We provide convenience methods to create standards configurations, for instance:
+    >>> from careamics.config import create_n2v_configuration
+    >>> config = create_n2v_configuration(
+    ...     experiment_name="n2v_experiment",
+    ...     data_type="array",
+    ...     axes="YX",
+    ...     patch_size=[64, 64],
+    ...     batch_size=32,
+    ...     num_epochs=100
+    ... )
+
+    The configuration can be exported to a dictionary using the model_dump method:
+    >>> config_dict = config.model_dump()
+
+    Configurations can also be exported or imported from yaml files:
+    >>> from careamics.config import save_configuration, load_configuration
+    >>> path_to_config = save_configuration(config, my_path / "config.yml")
+    >>> other_config = load_configuration(path_to_config)
+
+    Examples
+    --------
+    Minimum example:
+    >>> from careamics import configuration_factory
+    >>> config_dict = {
+    ...         "experiment_name": "N2V_experiment",
+    ...         "algorithm_config": {
+    ...             "algorithm": "n2v",
+    ...             "loss": "n2v",
+    ...             "model": {
+    ...                 "architecture": "UNet",
+    ...             },
+    ...         },
+    ...         "training_config": {
+    ...             "num_epochs": 200,
+    ...         },
+    ...         "data_config": {
+    ...             "data_type": "tiff",
+    ...             "patch_size": [64, 64],
+    ...             "axes": "SYX",
+    ...         },
+    ...     }
+    >>> config = configuration_factory(config_dict)
+    """
+
+    model_config = ConfigDict(
+        validate_assignment=True,
+        arbitrary_types_allowed=True,
+    )
+
+    # version
+    version: Literal["0.1.0"] = "0.1.0"
+    """CAREamics configuration version."""
+
+    # required parameters
+    experiment_name: str
+    """Name of the experiment, used to name logs and checkpoints."""
+
+    # Sub-configurations
+    algorithm_config: Union[UNetBasedAlgorithm, VAEBasedAlgorithm] = Field(
+        discriminator="algorithm"
+    )
+    """Algorithm configuration, holding all parameters required to configure the
+    model."""
+
+    data_config: GeneralDataConfig
+    """Data configuration, holding all parameters required to configure the training
+    data loader."""
+
+    training_config: TrainingConfig
+    """Training configuration, holding all parameters required to configure the
+    training process."""
+
+    @field_validator("experiment_name")
+    @classmethod
+    def no_symbol(cls, name: str) -> str:
+        """
+        Validate experiment name.
+
+        A valid experiment name is a non-empty string with only contains letters,
+        numbers, underscores, dashes and spaces.
+
+        Parameters
+        ----------
+        name : str
+            Name to validate.
+
+        Returns
+        -------
+        str
+            Validated name.
+
+        Raises
+        ------
+        ValueError
+            If the name is empty or contains invalid characters.
+        """
+        if len(name) == 0 or name.isspace():
+            raise ValueError("Experiment name is empty.")
+
+        # Validate using a regex that it contains only letters, numbers, underscores,
+        # dashes and spaces
+        if not re.match(r"^[a-zA-Z0-9_\- ]*$", name):
+            raise ValueError(
+                f"Experiment name contains invalid characters (got {name}). "
+                f"Only letters, numbers, underscores, dashes and spaces are allowed."
+            )
+
+        return name
+
+    @model_validator(mode="after")
+    def validate_3D(self: Self) -> Self:
+        """
+        Change algorithm dimensions to match data.axes.
+
+        Returns
+        -------
+        Self
+            Validated configuration.
+        """
+        if "Z" in self.data_config.axes and not self.algorithm_config.model.is_3D():
+            # change algorithm to 3D
+            self.algorithm_config.model.set_3D(True)
+        elif "Z" not in self.data_config.axes and self.algorithm_config.model.is_3D():
+            # change algorithm to 2D
+            self.algorithm_config.model.set_3D(False)
+
+        return self
+
+    def __str__(self) -> str:
+        """
+        Pretty string reprensenting the configuration.
+
+        Returns
+        -------
+        str
+            Pretty string.
+        """
+        return pformat(self.model_dump())
+
+    def set_3D(self, is_3D: bool, axes: str, patch_size: list[int]) -> None:
+        """
+        Set 3D flag and axes.
+
+        Parameters
+        ----------
+        is_3D : bool
+            Whether the algorithm is 3D or not.
+        axes : str
+            Axes of the data.
+        patch_size : list[int]
+            Patch size.
+        """
+        # set the flag and axes (this will not trigger validation at the config level)
+        self.algorithm_config.model.set_3D(is_3D)
+        self.data_config.set_3D(axes, patch_size)
+
+        # cheap hack: trigger validation
+        self.algorithm_config = self.algorithm_config
+
+    def get_algorithm_friendly_name(self) -> str:
+        """
+        Get the algorithm name.
+
+        Returns
+        -------
+        str
+            Algorithm name.
+        """
+        raise ValueError("Unknown algorithm.")
+
+    def get_algorithm_description(self) -> str:
+        """
+        Return a description of the algorithm.
+
+        This method is used to generate the README of the BioImage Model Zoo export.
+
+        Returns
+        -------
+        str
+            Description of the algorithm.
+        """
+        raise ValueError("No algorithm description available.")
+
+    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.
+        """
+        raise ValueError("No algorithm citations available.")
+
+    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.
+        """
+        raise ValueError("No algorithm references available.")
+
+    def get_algorithm_keywords(self) -> list[str]:
+        """
+        Get algorithm keywords.
+
+        Returns
+        -------
+        list[str]
+            List of keywords.
+        """
+        return ["CAREamics"]
+
+    def model_dump(
+        self,
+        *,
+        mode: Literal["json", "python"] | str = "python",
+        include: Any | None = None,
+        exclude: Any | None = None,
+        context: Any | None = None,
+        by_alias: bool = False,
+        exclude_unset: bool = False,
+        exclude_defaults: bool = False,
+        exclude_none: bool = True,
+        round_trip: bool = False,
+        warnings: bool | Literal["none", "warn", "error"] = True,
+        serialize_as_any: bool = False,
+    ) -> dict:
+        """
+        Override model_dump method in order to set default values.
+
+        As opposed to the parent model_dump method, this method sets exclude none by
+        default.
+
+        Parameters
+        ----------
+        mode : Literal['json', 'python'] | str, default='python'
+            The serialization format.
+        include : Any | None, default=None
+            Attributes to include.
+        exclude : Any | None, default=None
+            Attributes to exclude.
+        context : Any | None, default=None
+            Additional context to pass to the serialization functions.
+        by_alias : bool, default=False
+            Whether to use attribute aliases.
+        exclude_unset : bool, default=False
+            Whether to exclude fields that are not set.
+        exclude_defaults : bool, default=False
+            Whether to exclude fields that have default values.
+        exclude_none : bool, default=true
+            Whether to exclude fields that have None values.
+        round_trip : bool, default=False
+            Whether to dump and load the data to ensure that the output is a valid
+            representation.
+        warnings : bool | Literal['none', 'warn', 'error'], default=True
+            Whether to emit warnings.
+        serialize_as_any : bool, default=False
+            Whether to serialize all types as Any.
+
+        Returns
+        -------
+        dict
+            Dictionary containing the model parameters.
+        """
+        dictionary = super().model_dump(
+            mode=mode,
+            include=include,
+            exclude=exclude,
+            context=context,
+            by_alias=by_alias,
+            exclude_unset=exclude_unset,
+            exclude_defaults=exclude_defaults,
+            exclude_none=exclude_none,
+            round_trip=round_trip,
+            warnings=warnings,
+            serialize_as_any=serialize_as_any,
+        )
+
+        return dictionary

careamics/config/{configuration_factory.py → configuration_factories.py}

@@ -2,26 +2,93 @@
 
 from typing import Any, Literal, Optional, Union
 
-from .architectures import UNetModel
-from .configuration_model import Configuration
-from .data_model import DataConfig
-from .fcn_algorithm_model import FCNAlgorithmConfig
-from .support import (
+from pydantic import TypeAdapter
+
+from careamics.config.algorithms import CAREAlgorithm, N2NAlgorithm, N2VAlgorithm
+from careamics.config.architectures import UNetModel
+from careamics.config.care_configuration import CAREConfiguration
+from careamics.config.configuration import Configuration
+from careamics.config.data import DataConfig, N2VDataConfig
+from careamics.config.n2n_configuration import N2NConfiguration
+from careamics.config.n2v_configuration import N2VConfiguration
+from careamics.config.support import (
     SupportedArchitecture,
     SupportedPixelManipulation,
     SupportedTransform,
 )
-from .training_model import TrainingConfig
-from .transformations import (
+from careamics.config.training_model import TrainingConfig
+from careamics.config.transformations import (
+    N2V_TRANSFORMS_UNION,
+    SPATIAL_TRANSFORMS_UNION,
     N2VManipulateModel,
     XYFlipModel,
     XYRandomRotate90Model,
 )
 
 
-def _list_augmentations(
-    augmentations: Optional[list[Union[XYFlipModel, XYRandomRotate90Model]]],
-) -> list[Union[XYFlipModel, XYRandomRotate90Model]]:
+def configuration_factory(
+    configuration: dict[str, Any]
+) -> Union[N2VConfiguration, N2NConfiguration, CAREConfiguration]:
+    """
+    Create a configuration for training CAREamics.
+
+    Parameters
+    ----------
+    configuration : dict
+        Configuration dictionary.
+
+    Returns
+    -------
+    N2VConfiguration or N2NConfiguration or CAREConfiguration
+        Configuration for training CAREamics.
+    """
+    adapter: TypeAdapter = TypeAdapter(
+        Union[N2VConfiguration, N2NConfiguration, CAREConfiguration]
+    )
+    return adapter.validate_python(configuration)
+
+
+def algorithm_factory(
+    algorithm: dict[str, Any]
+) -> Union[N2VAlgorithm, N2NAlgorithm, CAREAlgorithm]:
+    """
+    Create an algorithm model for training CAREamics.
+
+    Parameters
+    ----------
+    algorithm : dict
+        Algorithm dictionary.
+
+    Returns
+    -------
+    N2VAlgorithm or N2NAlgorithm or CAREAlgorithm
+        Algorithm model for training CAREamics.
+    """
+    adapter: TypeAdapter = TypeAdapter(Union[N2VAlgorithm, N2NAlgorithm, CAREAlgorithm])
+    return adapter.validate_python(algorithm)
+
+
+def data_factory(data: dict[str, Any]) -> Union[DataConfig, N2VDataConfig]:
+    """
+    Create a data model for training CAREamics.
+
+    Parameters
+    ----------
+    data : dict
+        Data dictionary.
+
+    Returns
+    -------
+    DataConfig or N2VDataConfig
+        Data model for training CAREamics.
+    """
+    adapter: TypeAdapter = TypeAdapter(Union[DataConfig, N2VDataConfig])
+    return adapter.validate_python(data)
+
+
+def _list_spatial_augmentations(
+    augmentations: Optional[list[SPATIAL_TRANSFORMS_UNION]],
+) -> list[SPATIAL_TRANSFORMS_UNION]:
     """
     List the augmentations to apply.
 
@@ -44,7 +111,7 @@ def _list_augmentations(
         If there are duplicate transforms.
     """
     if augmentations is None:
-        transform_list: list[Union[XYFlipModel, XYRandomRotate90Model]] = [
+        transform_list: list[SPATIAL_TRANSFORMS_UNION] = [
             XYFlipModel(),
             XYRandomRotate90Model(),
         ]
@@ -123,7 +190,7 @@ def _create_configuration(
     patch_size: list[int],
     batch_size: int,
     num_epochs: int,
-    augmentations: list[Union[XYFlipModel, XYRandomRotate90Model]],
+    augmentations: Union[list[N2V_TRANSFORMS_UNION], list[SPATIAL_TRANSFORMS_UNION]],
     independent_channels: bool,
     loss: Literal["n2v", "mae", "mse"],
     n_channels_in: int,
@@ -188,21 +255,21 @@ def _create_configuration(
     )
 
     # algorithm model
-    algorithm_config = FCNAlgorithmConfig(
-        algorithm=algorithm,
-        loss=loss,
-        model=unet_model,
-    )
+    algorithm_config = {
+        "algorithm": algorithm,
+        "loss": loss,
+        "model": unet_model,
+    }
 
     # data model
-    data = DataConfig(
-        data_type=data_type,
-        axes=axes,
-        patch_size=patch_size,
-        batch_size=batch_size,
-        transforms=augmentations,
-        dataloader_params=dataloader_params,
-    )
+    data = {
+        "data_type": data_type,
+        "axes": axes,
+        "patch_size": patch_size,
+        "batch_size": batch_size,
+        "transforms": augmentations,
+        "dataloader_params": dataloader_params,
+    }
 
     # training model
     training = TrainingConfig(
@@ -212,14 +279,14 @@ def _create_configuration(
     )
 
     # create configuration
-    configuration = Configuration(
-        experiment_name=experiment_name,
-        algorithm_config=algorithm_config,
-        data_config=data,
-        training_config=training,
-    )
+    configuration = {
+        "experiment_name": experiment_name,
+        "algorithm_config": algorithm_config,
+        "data_config": data,
+        "training_config": training,
+    }
 
-    return configuration
+    return configuration_factory(configuration)
 
 
 # TODO reconsider naming once we officially support LVAE approaches
@@ -306,7 +373,7 @@ def _create_supervised_configuration(
         n_channels_out = n_channels_in
 
     # augmentations
-    transform_list = _list_augmentations(augmentations)
+    spatial_transform_list = _list_spatial_augmentations(augmentations)
 
     return _create_configuration(
         algorithm=algorithm,
@@ -316,7 +383,7 @@ def _create_supervised_configuration(
         patch_size=patch_size,
         batch_size=batch_size,
         num_epochs=num_epochs,
-        augmentations=transform_list,
+        augmentations=spatial_transform_list,
         independent_channels=independent_channels,
         loss=loss,
         n_channels_in=n_channels_in,
@@ -853,7 +920,7 @@ def create_n2v_configuration(
         n_channels = 1
 
     # augmentations
-    transform_list = _list_augmentations(augmentations)
+    spatial_transforms = _list_spatial_augmentations(augmentations)
 
     # create the N2VManipulate transform using the supplied parameters
     n2v_transform = N2VManipulateModel(
@@ -868,7 +935,7 @@ def create_n2v_configuration(
         struct_mask_axis=struct_n2v_axis,
         struct_mask_span=struct_n2v_span,
     )
-    transform_list.append(n2v_transform)
+    transform_list: list[N2V_TRANSFORMS_UNION] = spatial_transforms + [n2v_transform]
 
     return _create_configuration(
         algorithm="n2v",

careamics/config/configuration_io.py

@@ -0,0 +1,85 @@
+"""I/O functions for Configuration objects."""
+
+from pathlib import Path
+from typing import Union
+
+import yaml
+
+from careamics.config import Configuration, configuration_factory
+
+
+def load_configuration(path: Union[str, Path]) -> Configuration:
+    """
+    Load configuration from a yaml file.
+
+    Parameters
+    ----------
+    path : str or Path
+        Path to the configuration.
+
+    Returns
+    -------
+    Configuration
+        Configuration.
+
+    Raises
+    ------
+    FileNotFoundError
+        If the configuration file does not exist.
+    """
+    # load dictionary from yaml
+    if not Path(path).exists():
+        raise FileNotFoundError(
+            f"Configuration file {path} does not exist in " f" {Path.cwd()!s}"
+        )
+
+    dictionary = yaml.load(Path(path).open("r"), Loader=yaml.SafeLoader)
+
+    return configuration_factory(dictionary)
+
+
+def save_configuration(config: Configuration, path: Union[str, Path]) -> Path:
+    """
+    Save configuration to path.
+
+    Parameters
+    ----------
+    config : Configuration
+        Configuration to save.
+    path : str or Path
+        Path to a existing folder in which to save the configuration, or to a valid
+        configuration file path (uses a .yml or .yaml extension).
+
+    Returns
+    -------
+    Path
+        Path object representing the configuration.
+
+    Raises
+    ------
+    ValueError
+        If the path does not point to an existing directory or .yml file.
+    """
+    # make sure path is a Path object
+    config_path = Path(path)
+
+    # check if path is pointing to an existing directory or .yml file
+    if config_path.exists():
+        if config_path.is_dir():
+            config_path = Path(config_path, "config.yml")
+        elif config_path.suffix != ".yml" and config_path.suffix != ".yaml":
+            raise ValueError(
+                f"Path must be a directory or .yml or .yaml file (got {config_path})."
+            )
+    else:
+        if config_path.suffix != ".yml" and config_path.suffix != ".yaml":
+            raise ValueError(
+                f"Path must be a directory or .yml or .yaml file (got {config_path})."
+            )
+
+    # save configuration as dictionary to yaml
+    with open(config_path, "w") as f:
+        # dump configuration
+        yaml.dump(config.model_dump(), f, default_flow_style=False, sort_keys=False)
+
+    return config_path

careamics/config/data/__init__.py

@@ -0,0 +1,10 @@
+"""Data Pydantic configuration models."""
+
+__all__ = [
+    "DataConfig",
+    "GeneralDataConfig",
+    "N2VDataConfig",
+]
+
+from .data_model import DataConfig, GeneralDataConfig
+from .n2v_data_model import N2VDataConfig