careamics 0.0.9__py3-none-any.whl → 0.0.10__py3-none-any.whl

careamics/config/configuration_factories.py

@@ -2,76 +2,25 @@
 
 from typing import Annotated, Any, Literal, Optional, Union
 
-from pydantic import Discriminator, Tag, TypeAdapter
+from pydantic import Field, 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.data import DataConfig
 from careamics.config.support import (
-    SupportedAlgorithm,
     SupportedArchitecture,
     SupportedPixelManipulation,
     SupportedTransform,
 )
 from careamics.config.training_model import TrainingConfig
 from careamics.config.transformations import (
-    N2V_TRANSFORMS_UNION,
     SPATIAL_TRANSFORMS_UNION,
     N2VManipulateModel,
     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)
+from .configuration import Configuration
 
 
 def algorithm_factory(
@@ -90,28 +39,15 @@ def algorithm_factory(
     N2VAlgorithm or N2NAlgorithm or CAREAlgorithm
         Algorithm model for training CAREamics.
     """
-    adapter: TypeAdapter = TypeAdapter(Union[N2VAlgorithm, N2NAlgorithm, CAREAlgorithm])
+    adapter: TypeAdapter = TypeAdapter(
+        Annotated[
+            Union[N2VAlgorithm, N2NAlgorithm, CAREAlgorithm],
+            Field(discriminator="algorithm"),
+        ]
+    )
     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]:
@@ -208,70 +144,42 @@ def _create_unet_configuration(
     )
 
 
-def _create_configuration(
-    algorithm: Literal["n2v", "care", "n2n"],
-    experiment_name: str,
-    data_type: Literal["array", "tiff", "custom"],
+def _create_algorithm_configuration(
     axes: str,
-    patch_size: list[int],
-    batch_size: int,
-    num_epochs: int,
-    augmentations: Union[list[N2V_TRANSFORMS_UNION], list[SPATIAL_TRANSFORMS_UNION]],
-    independent_channels: bool,
+    algorithm: Literal["n2v", "care", "n2n"],
     loss: Literal["n2v", "mae", "mse"],
+    independent_channels: bool,
     n_channels_in: int,
     n_channels_out: int,
-    logger: Literal["wandb", "tensorboard", "none"],
     use_n2v2: bool = False,
     model_params: Optional[dict] = None,
-    train_dataloader_params: Optional[dict[str, Any]] = None,
-    val_dataloader_params: Optional[dict[str, Any]] = None,
-) -> Configuration:
+) -> dict:
     """
-    Create a configuration for training N2V, CARE or Noise2Noise.
+    Create a dictionary with the parameters of the algorithm model.
 
     Parameters
     ----------
+    axes : str
+        Axes of the data.
     algorithm : {"n2v", "care", "n2n"}
         Algorithm to use.
-    experiment_name : str
-        Name of the experiment.
-    data_type : {"array", "tiff", "custom"}
-        Type of the data.
-    axes : str
-        Axes of the data (e.g. SYX).
-    patch_size : list of int
-        Size of the patches along the spatial dimensions (e.g. [64, 64]).
-    batch_size : int
-        Batch size.
-    num_epochs : int
-        Number of epochs.
-    augmentations : list of transforms
-        List of transforms to apply, either both or one of XYFlipModel and
-        XYRandomRotate90Model.
-    independent_channels : bool
-        Whether to train all channels independently.
     loss : {"n2v", "mae", "mse"}
         Loss function to use.
+    independent_channels : bool
+        Whether to train all channels independently.
     n_channels_in : int
-        Number of channels in.
+        Number of input channels.
     n_channels_out : int
-        Number of channels out.
-    logger : {"wandb", "tensorboard", "none"}
-        Logger to use.
+        Number of output channels.
     use_n2v2 : bool, optional
         Whether to use N2V2, by default False.
     model_params : dict
         UNetModel parameters.
-    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
     -------
-    Configuration
-        Configuration for training N2V, CARE or Noise2Noise.
+    dict
+        Algorithm model as dictionnary with the specified parameters.
     """
     # model
     unet_model = _create_unet_configuration(
@@ -283,13 +191,47 @@ def _create_configuration(
         model_params=model_params,
     )
 
-    # algorithm model
-    algorithm_config = {
+    return {
         "algorithm": algorithm,
         "loss": loss,
         "model": unet_model,
     }
 
+
+def _create_data_configuration(
+    data_type: Literal["array", "tiff", "custom"],
+    axes: str,
+    patch_size: list[int],
+    batch_size: int,
+    augmentations: Union[list[SPATIAL_TRANSFORMS_UNION]],
+    train_dataloader_params: Optional[dict[str, Any]] = None,
+    val_dataloader_params: Optional[dict[str, Any]] = None,
+) -> DataConfig:
+    """
+    Create a dictionary with the parameters of the data model.
+
+    Parameters
+    ----------
+    data_type : {"array", "tiff", "custom"}
+        Type of the data.
+    axes : str
+        Axes of the data.
+    patch_size : list of int
+        Size of the patches along the spatial dimensions.
+    batch_size : int
+        Batch size.
+    augmentations : list of transforms
+        List of transforms to apply.
+    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
+    -------
+    DataConfig
+        Data model with the specified parameters.
+    """
     # data model
     data = {
         "data_type": data_type,
@@ -300,30 +242,44 @@ def _create_configuration(
     }
     # Don't override defaults set in DataConfig class
     if train_dataloader_params is not None:
+        # DataConfig enforces the presence of `shuffle` key in the dataloader parameters
+        if "shuffle" not in train_dataloader_params:
+            train_dataloader_params["shuffle"] = True
+
         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(
+    return DataConfig(**data)
+
+
+def _create_training_configuration(
+    num_epochs: int, logger: Literal["wandb", "tensorboard", "none"]
+) -> TrainingConfig:
+    """
+    Create a dictionary with the parameters of the training model.
+
+    Parameters
+    ----------
+    num_epochs : int
+        Number of epochs.
+    logger : {"wandb", "tensorboard", "none"}
+        Logger to use.
+
+    Returns
+    -------
+    TrainingConfig
+        Training model with the specified parameters.
+    """
+    return TrainingConfig(
         num_epochs=num_epochs,
-        batch_size=batch_size,
         logger=None if logger == "none" else logger,
     )
 
-    # create configuration
-    configuration = {
-        "experiment_name": experiment_name,
-        "algorithm_config": algorithm_config,
-        "data_config": data,
-        "training_config": training,
-    }
-
-    return configuration_factory(configuration)
-
 
 # TODO reconsider naming once we officially support LVAE approaches
-def _create_supervised_configuration(
+def _create_supervised_config_dict(
     algorithm: Literal["care", "n2n"],
     experiment_name: str,
     data_type: Literal["array", "tiff", "custom"],
@@ -331,7 +287,7 @@ def _create_supervised_configuration(
     patch_size: list[int],
     batch_size: int,
     num_epochs: int,
-    augmentations: Optional[list[Union[XYFlipModel, XYRandomRotate90Model]]] = None,
+    augmentations: Optional[list[SPATIAL_TRANSFORMS_UNION]] = None,
     independent_channels: bool = True,
     loss: Literal["mae", "mse"] = "mae",
     n_channels_in: Optional[int] = None,
@@ -340,7 +296,7 @@ def _create_supervised_configuration(
     model_params: Optional[dict] = None,
     train_dataloader_params: Optional[dict[str, Any]] = None,
     val_dataloader_params: Optional[dict[str, Any]] = None,
-) -> Configuration:
+) -> dict:
     """
     Create a configuration for training CARE or Noise2Noise.
 
@@ -411,25 +367,41 @@ def _create_supervised_configuration(
     # augmentations
     spatial_transform_list = _list_spatial_augmentations(augmentations)
 
-    return _create_configuration(
+    # algorithm
+    algorithm_params = _create_algorithm_configuration(
+        axes=axes,
         algorithm=algorithm,
-        experiment_name=experiment_name,
+        loss=loss,
+        independent_channels=independent_channels,
+        n_channels_in=n_channels_in,
+        n_channels_out=n_channels_out,
+        model_params=model_params,
+    )
+
+    # data
+    data_params = _create_data_configuration(
         data_type=data_type,
         axes=axes,
         patch_size=patch_size,
         batch_size=batch_size,
-        num_epochs=num_epochs,
         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,
         train_dataloader_params=train_dataloader_params,
         val_dataloader_params=val_dataloader_params,
     )
 
+    # training
+    training_params = _create_training_configuration(
+        num_epochs=num_epochs,
+        logger=logger,
+    )
+
+    return {
+        "experiment_name": experiment_name,
+        "algorithm_config": algorithm_params,
+        "data_config": data_params,
+        "training_config": training_params,
+    }
+
 
 def create_care_configuration(
     experiment_name: str,
@@ -580,23 +552,25 @@ def create_care_configuration(
     ...     n_channels_out=1 # if applicable
     ... )
     """
-    return _create_supervised_configuration(
-        algorithm="care",
-        experiment_name=experiment_name,
-        data_type=data_type,
-        axes=axes,
-        patch_size=patch_size,
-        batch_size=batch_size,
-        num_epochs=num_epochs,
-        augmentations=augmentations,
-        independent_channels=independent_channels,
-        loss=loss,
-        n_channels_in=n_channels_in,
-        n_channels_out=n_channels_out,
-        logger=logger,
-        model_params=model_params,
-        train_dataloader_params=train_dataloader_params,
-        val_dataloader_params=val_dataloader_params,
+    return Configuration(
+        **_create_supervised_config_dict(
+            algorithm="care",
+            experiment_name=experiment_name,
+            data_type=data_type,
+            axes=axes,
+            patch_size=patch_size,
+            batch_size=batch_size,
+            num_epochs=num_epochs,
+            augmentations=augmentations,
+            independent_channels=independent_channels,
+            loss=loss,
+            n_channels_in=n_channels_in,
+            n_channels_out=n_channels_out,
+            logger=logger,
+            model_params=model_params,
+            train_dataloader_params=train_dataloader_params,
+            val_dataloader_params=val_dataloader_params,
+        )
     )
 
 
@@ -749,23 +723,25 @@ def create_n2n_configuration(
     ...     n_channels_out=1 # if applicable
     ... )
     """
-    return _create_supervised_configuration(
-        algorithm="n2n",
-        experiment_name=experiment_name,
-        data_type=data_type,
-        axes=axes,
-        patch_size=patch_size,
-        batch_size=batch_size,
-        num_epochs=num_epochs,
-        augmentations=augmentations,
-        independent_channels=independent_channels,
-        loss=loss,
-        n_channels_in=n_channels_in,
-        n_channels_out=n_channels_out,
-        logger=logger,
-        model_params=model_params,
-        train_dataloader_params=train_dataloader_params,
-        val_dataloader_params=val_dataloader_params,
+    return Configuration(
+        **_create_supervised_config_dict(
+            algorithm="n2n",
+            experiment_name=experiment_name,
+            data_type=data_type,
+            axes=axes,
+            patch_size=patch_size,
+            batch_size=batch_size,
+            num_epochs=num_epochs,
+            augmentations=augmentations,
+            independent_channels=independent_channels,
+            loss=loss,
+            n_channels_in=n_channels_in,
+            n_channels_out=n_channels_out,
+            logger=logger,
+            model_params=model_params,
+            train_dataloader_params=train_dataloader_params,
+            val_dataloader_params=val_dataloader_params,
+        )
     )
 
 
@@ -995,24 +971,40 @@ def create_n2v_configuration(
         struct_mask_axis=struct_n2v_axis,
         struct_mask_span=struct_n2v_span,
     )
-    transform_list: list[N2V_TRANSFORMS_UNION] = spatial_transforms + [n2v_transform]
 
-    return _create_configuration(
-        algorithm="n2v",
-        experiment_name=experiment_name,
-        data_type=data_type,
+    # algorithm
+    algorithm_params = _create_algorithm_configuration(
         axes=axes,
-        patch_size=patch_size,
-        batch_size=batch_size,
-        num_epochs=num_epochs,
-        augmentations=transform_list,
-        independent_channels=independent_channels,
+        algorithm="n2v",
         loss="n2v",
-        use_n2v2=use_n2v2,
+        independent_channels=independent_channels,
         n_channels_in=n_channels,
         n_channels_out=n_channels,
-        logger=logger,
+        use_n2v2=use_n2v2,
         model_params=model_params,
+    )
+    algorithm_params["n2v_config"] = n2v_transform
+
+    # data
+    data_params = _create_data_configuration(
+        data_type=data_type,
+        axes=axes,
+        patch_size=patch_size,
+        batch_size=batch_size,
+        augmentations=spatial_transforms,
         train_dataloader_params=train_dataloader_params,
         val_dataloader_params=val_dataloader_params,
     )
+
+    # training
+    training_params = _create_training_configuration(
+        num_epochs=num_epochs,
+        logger=logger,
+    )
+
+    return Configuration(
+        experiment_name=experiment_name,
+        algorithm_config=algorithm_params,
+        data_config=data_params,
+        training_config=training_params,
+    )

careamics/config/configuration_io.py

@@ -5,7 +5,7 @@ from typing import Union
 
 import yaml
 
-from careamics.config import Configuration, configuration_factory
+from careamics.config import Configuration
 
 
 def load_configuration(path: Union[str, Path]) -> Configuration:
@@ -35,7 +35,7 @@ def load_configuration(path: Union[str, Path]) -> Configuration:
 
     dictionary = yaml.load(Path(path).open("r"), Loader=yaml.SafeLoader)
 
-    return configuration_factory(dictionary)
+    return Configuration(**dictionary)
 
 
 def save_configuration(config: Configuration, path: Union[str, Path]) -> Path:

careamics/config/data/__init__.py

@@ -2,9 +2,6 @@
 
 __all__ = [
     "DataConfig",
-    "GeneralDataConfig",
-    "N2VDataConfig",
 ]
 
-from .data_model import DataConfig, GeneralDataConfig
-from .n2v_data_model import N2VDataConfig
+from .data_model import DataConfig

careamics/config/data/data_model.py

@@ -19,7 +19,7 @@ from pydantic import (
 )
 from typing_extensions import Self
 
-from ..transformations import N2V_TRANSFORMS_UNION, XYFlipModel, XYRandomRotate90Model
+from ..transformations import XYFlipModel, XYRandomRotate90Model
 from ..validators import check_axes_validity, patch_size_ge_than_8_power_of_2
 
 
@@ -46,8 +46,46 @@ Float = Annotated[float, PlainSerializer(np_float_to_scientific_str, return_type
 """Annotated float type, used to serialize floats to strings."""
 
 
-class GeneralDataConfig(BaseModel):
-    """General data configuration."""
+class DataConfig(BaseModel):
+    """Data configuration.
+
+    If std is specified, mean must be specified as well. Note that setting the std first
+    and then the mean (if they were both `None` before) will raise a validation error.
+    Prefer instead `set_mean_and_std` to set both at once. Means and stds are expected
+    to be lists of floats, one for each channel. For supervised tasks, the mean and std
+    of the target could be different from the input data.
+
+    All supported transforms are defined in the SupportedTransform enum.
+
+    Examples
+    --------
+    Minimum example:
+
+    >>> data = DataConfig(
+    ...     data_type="array", # defined in SupportedData
+    ...     patch_size=[128, 128],
+    ...     batch_size=4,
+    ...     axes="YX"
+    ... )
+
+    To change the image_means and image_stds of the data:
+    >>> data.set_means_and_stds(image_means=[214.3], image_stds=[84.5])
+
+    One can pass also a list of transformations, by keyword, using the
+    SupportedTransform value:
+    >>> from careamics.config.support import SupportedTransform
+    >>> data = DataConfig(
+    ...     data_type="tiff",
+    ...     patch_size=[128, 128],
+    ...     batch_size=4,
+    ...     axes="YX",
+    ...     transforms=[
+    ...         {
+    ...             "name": "XYFlip",
+    ...         }
+    ...     ]
+    ... )
+    """
 
     # Pydantic class configuration
     model_config = ConfigDict(
@@ -88,10 +126,7 @@ class GeneralDataConfig(BaseModel):
     """Standard deviations of the target data across channels, used for
     normalization."""
 
-    # defining as Sequence allows assigning subclasses of TransformModel without mypy
-    # complaining, this is important for instance to differentiate N2VDataConfig and
-    # DataConfig
-    transforms: Sequence[N2V_TRANSFORMS_UNION] = Field(
+    transforms: Sequence[Union[XYFlipModel, XYRandomRotate90Model]] = Field(
         default=[
             XYFlipModel(),
             XYRandomRotate90Model(),
@@ -104,7 +139,9 @@ class GeneralDataConfig(BaseModel):
     train_dataloader_params: dict[str, Any] = Field(
         default={"shuffle": True}, validate_default=True
     )
-    """Dictionary of PyTorch training dataloader parameters."""
+    """Dictionary of PyTorch training dataloader parameters. The dataloader parameters,
+    should include the `shuffle` key, which is set to `True` by default. We strongly
+    recommend to keep it as `True` to ensure the best training results."""
 
     val_dataloader_params: dict[str, Any] = Field(default={})
     """Dictionary of PyTorch validation dataloader parameters."""
@@ -207,7 +244,7 @@ class GeneralDataConfig(BaseModel):
         ):
             warn(
                 "Dataloader parameters include `shuffle=False`, this will be passed to "
-                "the training dataloader and may result in bad results.",
+                "the training dataloader and may lead to lower quality results.",
                 stacklevel=1,
             )
         return train_dataloader_params
@@ -363,56 +400,3 @@ class GeneralDataConfig(BaseModel):
             Patch size.
         """
         self._update(axes=axes, patch_size=patch_size)
-
-
-class DataConfig(GeneralDataConfig):
-    """
-    Data configuration.
-
-    If std is specified, mean must be specified as well. Note that setting the std first
-    and then the mean (if they were both `None` before) will raise a validation error.
-    Prefer instead `set_mean_and_std` to set both at once. Means and stds are expected
-    to be lists of floats, one for each channel. For supervised tasks, the mean and std
-    of the target could be different from the input data.
-
-    All supported transforms are defined in the SupportedTransform enum.
-
-    Examples
-    --------
-    Minimum example:
-
-    >>> data = DataConfig(
-    ...     data_type="array", # defined in SupportedData
-    ...     patch_size=[128, 128],
-    ...     batch_size=4,
-    ...     axes="YX"
-    ... )
-
-    To change the image_means and image_stds of the data:
-    >>> data.set_means_and_stds(image_means=[214.3], image_stds=[84.5])
-
-    One can pass also a list of transformations, by keyword, using the
-    SupportedTransform value:
-    >>> from careamics.config.support import SupportedTransform
-    >>> data = DataConfig(
-    ...     data_type="tiff",
-    ...     patch_size=[128, 128],
-    ...     batch_size=4,
-    ...     axes="YX",
-    ...     transforms=[
-    ...         {
-    ...             "name": "XYFlip",
-    ...         }
-    ...     ]
-    ... )
-    """
-
-    transforms: Sequence[Union[XYFlipModel, XYRandomRotate90Model]] = Field(
-        default=[
-            XYFlipModel(),
-            XYRandomRotate90Model(),
-        ],
-        validate_default=True,
-    )
-    """List of transformations to apply to the data, available transforms are defined
-    in SupportedTransform. This excludes N2V specific transformations."""

careamics/config/support/supported_transforms.py

@@ -8,5 +8,5 @@ class SupportedTransform(str, BaseEnum):
 
     XY_FLIP = "XYFlip"
     XY_RANDOM_ROTATE90 = "XYRandomRotate90"
-    N2V_MANIPULATE = "N2VManipulate"
     NORMALIZE = "Normalize"
+    N2V_MANIPULATE = "N2VManipulate"