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

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

@@ -1,27 +1,120 @@
 """Convenience functions to create configurations for training and inference."""
 
-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 typing import Annotated, Any, Literal, Optional, Union
+
+from pydantic import Discriminator, Tag, 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 (
+    SupportedAlgorithm,
     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 _algorithm_config_discriminator(value: Union[dict, Configuration]) -> str:
+    """Discriminate algorithm-specific configurations based on the algorithm.
+
+    Parameters
+    ----------
+    value : Any
+        Value to discriminate.
+
+    Returns
+    -------
+    str
+        Discriminator value.
+    """
+    if isinstance(value, dict):
+        return value["algorithm_config"]["algorithm"]
+    return value.algorithm_config.algorithm
+
+
+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(
+        Annotated[
+            Union[
+                Annotated[N2VConfiguration, Tag(SupportedAlgorithm.N2V.value)],
+                Annotated[N2NConfiguration, Tag(SupportedAlgorithm.N2N.value)],
+                Annotated[CAREConfiguration, Tag(SupportedAlgorithm.CARE.value)],
+            ],
+            Discriminator(_algorithm_config_discriminator),
+        ]
+    )
+    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 +137,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 +216,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,
@@ -131,7 +224,8 @@ def _create_configuration(
     logger: Literal["wandb", "tensorboard", "none"],
     use_n2v2: bool = False,
     model_params: Optional[dict] = None,
-    dataloader_params: Optional[dict] = None,
+    train_dataloader_params: Optional[dict[str, Any]] = None,
+    val_dataloader_params: Optional[dict[str, Any]] = None,
 ) -> Configuration:
     """
     Create a configuration for training N2V, CARE or Noise2Noise.
@@ -169,8 +263,10 @@ def _create_configuration(
         Whether to use N2V2, by default False.
     model_params : dict
         UNetModel parameters.
-    dataloader_params : dict
-        Parameters for the dataloader, see PyTorch notes, by default None.
+    train_dataloader_params : dict
+        Parameters for the training dataloader, see PyTorch notes, by default None.
+    val_dataloader_params : dict
+        Parameters for the validation dataloader, see PyTorch notes, by default None.
 
     Returns
     -------
@@ -188,21 +284,25 @@ 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,
+    }
+    # Don't override defaults set in DataConfig class
+    if train_dataloader_params is not None:
+        data["train_dataloader_params"] = train_dataloader_params
+    if val_dataloader_params is not None:
+        data["val_dataloader_params"] = val_dataloader_params
 
     # training model
     training = TrainingConfig(
@@ -212,14 +312,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
@@ -238,7 +338,8 @@ def _create_supervised_configuration(
     n_channels_out: Optional[int] = None,
     logger: Literal["wandb", "tensorboard", "none"] = "none",
     model_params: Optional[dict] = None,
-    dataloader_params: Optional[dict] = None,
+    train_dataloader_params: Optional[dict[str, Any]] = None,
+    val_dataloader_params: Optional[dict[str, Any]] = None,
 ) -> Configuration:
     """
     Create a configuration for training CARE or Noise2Noise.
@@ -275,8 +376,10 @@ def _create_supervised_configuration(
         Logger to use, by default "none".
     model_params : dict, optional
         UNetModel parameters, by default {}.
-    dataloader_params : dict, optional
-        Parameters for the dataloader, see PyTorch notes, by default None.
+    train_dataloader_params : dict
+        Parameters for the training dataloader, see PyTorch notes, by default None.
+    val_dataloader_params : dict
+        Parameters for the validation dataloader, see PyTorch notes, by default None.
 
     Returns
     -------
@@ -306,7 +409,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,14 +419,15 @@ 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,
         n_channels_out=n_channels_out,
         logger=logger,
         model_params=model_params,
-        dataloader_params=dataloader_params,
+        train_dataloader_params=train_dataloader_params,
+        val_dataloader_params=val_dataloader_params,
     )
 
 
@@ -341,7 +445,8 @@ def create_care_configuration(
     n_channels_out: Optional[int] = None,
     logger: Literal["wandb", "tensorboard", "none"] = "none",
     model_params: Optional[dict] = None,
-    dataloader_params: Optional[dict] = None,
+    train_dataloader_params: Optional[dict[str, Any]] = None,
+    val_dataloader_params: Optional[dict[str, Any]] = None,
 ) -> Configuration:
     """
     Create a configuration for training CARE.
@@ -394,8 +499,14 @@ def create_care_configuration(
         Logger to use.
     model_params : dict, default=None
         UNetModel parameters.
-    dataloader_params : dict, optional
-        Parameters for the dataloader, see PyTorch notes, by default None.
+    train_dataloader_params : dict, optional
+        Parameters for the training dataloader, see the PyTorch docs for `DataLoader`.
+        If left as `None`, the dict `{"shuffle": True}` will be used, this is set in
+        the `GeneralDataConfig`.
+    val_dataloader_params : dict, optional
+        Parameters for the validation dataloader, see PyTorch the docs for `DataLoader`.
+        If left as `None`, the empty dict `{}` will be used, this is set in the
+        `GeneralDataConfig`.
 
     Returns
     -------
@@ -484,7 +595,8 @@ def create_care_configuration(
         n_channels_out=n_channels_out,
         logger=logger,
         model_params=model_params,
-        dataloader_params=dataloader_params,
+        train_dataloader_params=train_dataloader_params,
+        val_dataloader_params=val_dataloader_params,
     )
 
 
@@ -502,7 +614,8 @@ def create_n2n_configuration(
     n_channels_out: Optional[int] = None,
     logger: Literal["wandb", "tensorboard", "none"] = "none",
     model_params: Optional[dict] = None,
-    dataloader_params: Optional[dict] = None,
+    train_dataloader_params: Optional[dict[str, Any]] = None,
+    val_dataloader_params: Optional[dict[str, Any]] = None,
 ) -> Configuration:
     """
     Create a configuration for training Noise2Noise.
@@ -555,8 +668,14 @@ def create_n2n_configuration(
         Logger to use, by default "none".
     model_params : dict, optional
         UNetModel parameters, by default {}.
-    dataloader_params : dict, optional
-        Parameters for the dataloader, see PyTorch notes, by default None.
+    train_dataloader_params : dict, optional
+        Parameters for the training dataloader, see the PyTorch docs for `DataLoader`.
+        If left as `None`, the dict `{"shuffle": True}` will be used, this is set in
+        the `GeneralDataConfig`.
+    val_dataloader_params : dict, optional
+        Parameters for the validation dataloader, see PyTorch the docs for `DataLoader`.
+        If left as `None`, the empty dict `{}` will be used, this is set in the
+        `GeneralDataConfig`.
 
     Returns
     -------
@@ -645,7 +764,8 @@ def create_n2n_configuration(
         n_channels_out=n_channels_out,
         logger=logger,
         model_params=model_params,
-        dataloader_params=dataloader_params,
+        train_dataloader_params=train_dataloader_params,
+        val_dataloader_params=val_dataloader_params,
     )
 
 
@@ -666,7 +786,8 @@ def create_n2v_configuration(
     struct_n2v_span: int = 5,
     logger: Literal["wandb", "tensorboard", "none"] = "none",
     model_params: Optional[dict] = None,
-    dataloader_params: Optional[dict] = None,
+    train_dataloader_params: Optional[dict[str, Any]] = None,
+    val_dataloader_params: Optional[dict[str, Any]] = None,
 ) -> Configuration:
     """
     Create a configuration for training Noise2Void.
@@ -745,8 +866,14 @@ def create_n2v_configuration(
         Logger to use, by default "none".
     model_params : dict, optional
         UNetModel parameters, by default None.
-    dataloader_params : dict, optional
-        Parameters for the dataloader, see PyTorch notes, by default None.
+    train_dataloader_params : dict, optional
+        Parameters for the training dataloader, see the PyTorch docs for `DataLoader`.
+        If left as `None`, the dict `{"shuffle": True}` will be used, this is set in
+        the `GeneralDataConfig`.
+    val_dataloader_params : dict, optional
+        Parameters for the validation dataloader, see PyTorch the docs for `DataLoader`.
+        If left as `None`, the empty dict `{}` will be used, this is set in the
+        `GeneralDataConfig`.
 
     Returns
     -------
@@ -853,7 +980,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 +995,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",
@@ -886,5 +1013,6 @@ def create_n2v_configuration(
         n_channels_out=n_channels,
         logger=logger,
         model_params=model_params,
-        dataloader_params=dataloader_params,
+        train_dataloader_params=train_dataloader_params,
+        val_dataloader_params=val_dataloader_params,
     )

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