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

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
@@ -234,8 +301,8 @@ def _create_supervised_configuration(
     augmentations: Optional[list[Union[XYFlipModel, XYRandomRotate90Model]]] = None,
     independent_channels: bool = True,
     loss: Literal["mae", "mse"] = "mae",
-    n_channels_in: int = 1,
-    n_channels_out: int = 1,
+    n_channels_in: Optional[int] = None,
+    n_channels_out: Optional[int] = None,
     logger: Literal["wandb", "tensorboard", "none"] = "none",
     model_params: Optional[dict] = None,
     dataloader_params: Optional[dict] = None,
@@ -267,10 +334,10 @@ def _create_supervised_configuration(
         Whether to train all channels independently, by default False.
     loss : Literal["mae", "mse"], optional
         Loss function to use, by default "mae".
-    n_channels_in : int, optional
-        Number of channels in, by default 1.
-    n_channels_out : int, optional
-        Number of channels out, by default 1.
+    n_channels_in : int or None, default=None
+        Number of channels in.
+    n_channels_out : int or None, default=None
+        Number of channels out.
     logger : Literal["wandb", "tensorboard", "none"], optional
         Logger to use, by default "none".
     model_params : dict, optional
@@ -282,21 +349,31 @@ def _create_supervised_configuration(
     -------
     Configuration
         Configuration for training CARE or Noise2Noise.
+
+    Raises
+    ------
+    ValueError
+        If the number of channels is not specified when using channels.
+    ValueError
+        If the number of channels is specified but "C" is not in the axes.
     """
     # if there are channels, we need to specify their number
-    if "C" in axes and n_channels_in == 1:
-        raise ValueError(
-            f"Number of channels in must be specified when using channels "
-            f"(got {n_channels_in} channel)."
-        )
-    elif "C" not in axes and n_channels_in > 1:
+    if "C" in axes and n_channels_in is None:
+        raise ValueError("Number of channels in must be specified when using channels ")
+    elif "C" not in axes and (n_channels_in is not None and n_channels_in > 1):
         raise ValueError(
             f"C is not present in the axes, but number of channels is specified "
             f"(got {n_channels_in} channels)."
         )
 
+    if n_channels_in is None:
+        n_channels_in = 1
+
+    if n_channels_out is None:
+        n_channels_out = n_channels_in
+
     # augmentations
-    transform_list = _list_augmentations(augmentations)
+    spatial_transform_list = _list_spatial_augmentations(augmentations)
 
     return _create_configuration(
         algorithm=algorithm,
@@ -306,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,
@@ -327,8 +404,8 @@ def create_care_configuration(
     augmentations: Optional[list[Union[XYFlipModel, XYRandomRotate90Model]]] = None,
     independent_channels: bool = True,
     loss: Literal["mae", "mse"] = "mae",
-    n_channels_in: int = 1,
-    n_channels_out: int = -1,
+    n_channels_in: Optional[int] = None,
+    n_channels_out: Optional[int] = None,
     logger: Literal["wandb", "tensorboard", "none"] = "none",
     model_params: Optional[dict] = None,
     dataloader_params: Optional[dict] = None,
@@ -374,16 +451,16 @@ def create_care_configuration(
         and XYRandomRotate90 (in XY) to the images.
     independent_channels : bool, optional
         Whether to train all channels independently, by default False.
-    loss : Literal["mae", "mse"], optional
-        Loss function to use, by default "mae".
-    n_channels_in : int, optional
-        Number of channels in, by default 1.
-    n_channels_out : int, optional
-        Number of channels out, by default -1.
-    logger : Literal["wandb", "tensorboard", "none"], optional
-        Logger to use, by default "none".
-    model_params : dict, optional
-        UNetModel parameters, by default None.
+    loss : Literal["mae", "mse"], default="mae"
+        Loss function to use.
+    n_channels_in : int or None, default=None
+        Number of channels in.
+    n_channels_out : int or None, default=None
+        Number of channels out.
+    logger : Literal["wandb", "tensorboard", "none"], default="none"
+        Logger to use.
+    model_params : dict, default=None
+        UNetModel parameters.
     dataloader_params : dict, optional
         Parameters for the dataloader, see PyTorch notes, by default None.
 
@@ -459,9 +536,6 @@ def create_care_configuration(
     ...     n_channels_out=1 # if applicable
     ... )
     """
-    if n_channels_out == -1:
-        n_channels_out = n_channels_in
-
     return _create_supervised_configuration(
         algorithm="care",
         experiment_name=experiment_name,
@@ -491,8 +565,8 @@ def create_n2n_configuration(
     augmentations: Optional[list[Union[XYFlipModel, XYRandomRotate90Model]]] = None,
     independent_channels: bool = True,
     loss: Literal["mae", "mse"] = "mae",
-    n_channels_in: int = 1,
-    n_channels_out: int = -1,
+    n_channels_in: Optional[int] = None,
+    n_channels_out: Optional[int] = None,
     logger: Literal["wandb", "tensorboard", "none"] = "none",
     model_params: Optional[dict] = None,
     dataloader_params: Optional[dict] = None,
@@ -540,10 +614,10 @@ def create_n2n_configuration(
         Whether to train all channels independently, by default False.
     loss : Literal["mae", "mse"], optional
         Loss function to use, by default "mae".
-    n_channels_in : int, optional
-        Number of channels in, by default 1.
-    n_channels_out : int, optional
-        Number of channels out, by default -1.
+    n_channels_in : int or None, default=None
+        Number of channels in.
+    n_channels_out : int or None, default=None
+        Number of channels out.
     logger : Literal["wandb", "tensorboard", "none"], optional
         Logger to use, by default "none".
     model_params : dict, optional
@@ -623,9 +697,6 @@ def create_n2n_configuration(
     ...     n_channels_out=1 # if applicable
     ... )
     """
-    if n_channels_out == -1:
-        n_channels_out = n_channels_in
-
     return _create_supervised_configuration(
         algorithm="n2n",
         experiment_name=experiment_name,
@@ -655,7 +726,7 @@ def create_n2v_configuration(
     augmentations: Optional[list[Union[XYFlipModel, XYRandomRotate90Model]]] = None,
     independent_channels: bool = True,
     use_n2v2: bool = False,
-    n_channels: int = 1,
+    n_channels: Optional[int] = None,
     roi_size: int = 11,
     masked_pixel_percentage: float = 0.2,
     struct_n2v_axis: Literal["horizontal", "vertical", "none"] = "none",
@@ -727,8 +798,8 @@ def create_n2v_configuration(
         Whether to train all channels together, by default True.
     use_n2v2 : bool, optional
         Whether to use N2V2, by default False.
-    n_channels : int, optional
-        Number of channels (in and out), by default 1.
+    n_channels : int or None, default=None
+        Number of channels (in and out).
     roi_size : int, optional
         N2V pixel manipulation area, by default 11.
     masked_pixel_percentage : float, optional
@@ -837,19 +908,19 @@ def create_n2v_configuration(
     ... )
     """
     # if there are channels, we need to specify their number
-    if "C" in axes and n_channels == 1:
-        raise ValueError(
-            f"Number of channels must be specified when using channels "
-            f"(got {n_channels} channel)."
-        )
-    elif "C" not in axes and n_channels > 1:
+    if "C" in axes and n_channels is None:
+        raise ValueError("Number of channels must be specified when using channels.")
+    elif "C" not in axes and (n_channels is not None and n_channels > 1):
         raise ValueError(
             f"C is not present in the axes, but number of channels is specified "
             f"(got {n_channels} channel)."
         )
 
+    if n_channels is None:
+        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(
@@ -864,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