careamics 0.1.0rc3__py3-none-any.whl → 0.1.0rc5__py3-none-any.whl

careamics/config/configuration_factory.py

@@ -2,13 +2,11 @@
 
 from typing import Any, Dict, List, Literal, Optional, Tuple, Union
 
-from albumentations import Compose
-
-from .algorithm_model import AlgorithmModel
+from .algorithm_model import AlgorithmConfig
 from .architectures import UNetModel
 from .configuration_model import Configuration
-from .data_model import DataModel
-from .inference_model import InferenceModel
+from .data_model import DataConfig
+from .inference_model import InferenceConfig
 from .support import (
     SupportedAlgorithm,
     SupportedArchitecture,
@@ -16,10 +14,11 @@ from .support import (
     SupportedPixelManipulation,
     SupportedTransform,
 )
-from .training_model import TrainingModel
+from .training_model import TrainingConfig
 
 
-def create_n2n_configuration(
+def _create_supervised_configuration(
+    algorithm: Literal["care", "n2n"],
     experiment_name: str,
     data_type: Literal["array", "tiff", "custom"],
     axes: str,
@@ -27,28 +26,20 @@ def create_n2n_configuration(
     batch_size: int,
     num_epochs: int,
     use_augmentations: bool = True,
-    use_n2v2: bool = False,
-    n_channels: int = 1,
+    independent_channels: bool = False,
+    loss: Literal["mae", "mse"] = "mae",
+    n_channels_in: int = 1,
+    n_channels_out: int = 1,
     logger: Literal["wandb", "tensorboard", "none"] = "none",
     model_kwargs: Optional[dict] = None,
 ) -> Configuration:
     """
-    Create a configuration for training N2V.
-
-    If "Z" is present in `axes`, then `path_size` must be a list of length 3, otherwise
-    2.
-
-    By setting `use_augmentations` to False, the only transformation applied will be
-    normalization and N2V manipulation.
-
-    The parameter `use_n2v2` overrides the corresponding `n2v2` that can be passed
-    in `model_kwargs`.
-
-    If you pass "horizontal" or "vertical" to `struct_n2v_axis`, then structN2V mask
-    will be applied to each manipulated pixel.
+    Create a configuration for training CARE or Noise2Noise.
 
     Parameters
     ----------
+    algorithm : Literal["care", "n2n"]
+        Algorithm to use.
     experiment_name : str
         Name of the experiment.
     data_type : Literal["array", "tiff", "custom"]
@@ -63,18 +54,14 @@ def create_n2n_configuration(
         Number of epochs.
     use_augmentations : bool, optional
         Whether to use augmentations, 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.
-    roi_size : int, optional
-        N2V pixel manipulation area, by default 11.
-    masked_pixel_percentage : float, optional
-        Percentage of pixels masked in each patch, by default 0.2.
-    struct_n2v_axis : Literal["horizontal", "vertical", "none"], optional
-        Axis along which to apply structN2V mask, by default "none".
-    struct_n2v_span : int, optional
-        Span of the structN2V mask, by default 5.
+    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_kwargs : dict, optional
@@ -83,15 +70,27 @@ def create_n2n_configuration(
     Returns
     -------
     Configuration
-        Configuration for training N2V.
+        Configuration for training CARE or Noise2Noise.
     """
+    # 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:
+        raise ValueError(
+            f"C is not present in the axes, but number of channels is specified "
+            f"(got {n_channels_in} channels)."
+        )
+
     # model
     if model_kwargs is None:
         model_kwargs = {}
-    model_kwargs["n2v2"] = use_n2v2
     model_kwargs["conv_dims"] = 3 if "Z" in axes else 2
-    model_kwargs["in_channels"] = n_channels
-    model_kwargs["num_classes"] = n_channels
+    model_kwargs["in_channels"] = n_channels_in
+    model_kwargs["num_classes"] = n_channels_out
+    model_kwargs["independent_channels"] = independent_channels
 
     unet_model = UNetModel(
         architecture=SupportedArchitecture.UNET.value,
@@ -99,9 +98,9 @@ def create_n2n_configuration(
     )
 
     # algorithm model
-    algorithm = AlgorithmModel(
-        algorithm=SupportedAlgorithm.N2V.value,
-        loss=SupportedLoss.N2V.value,
+    algorithm = AlgorithmConfig(
+        algorithm=algorithm,
+        loss=loss,
         model=unet_model,
     )
 
@@ -126,7 +125,7 @@ def create_n2n_configuration(
         ]
 
     # data model
-    data = DataModel(
+    data = DataConfig(
         data_type=data_type,
         axes=axes,
         patch_size=patch_size,
@@ -135,7 +134,7 @@ def create_n2n_configuration(
     )
 
     # training model
-    training = TrainingModel(
+    training = TrainingConfig(
         num_epochs=num_epochs,
         batch_size=batch_size,
         logger=None if logger == "none" else logger,
@@ -152,6 +151,175 @@ def create_n2n_configuration(
     return configuration
 
 
+def create_care_configuration(
+    experiment_name: str,
+    data_type: Literal["array", "tiff", "custom"],
+    axes: str,
+    patch_size: List[int],
+    batch_size: int,
+    num_epochs: int,
+    use_augmentations: bool = True,
+    independent_channels: bool = False,
+    loss: Literal["mae", "mse"] = "mae",
+    n_channels_in: int = 1,
+    n_channels_out: int = -1,
+    logger: Literal["wandb", "tensorboard", "none"] = "none",
+    model_kwargs: Optional[dict] = None,
+) -> Configuration:
+    """
+    Create a configuration for training CARE.
+
+    If "Z" is present in `axes`, then `path_size` must be a list of length 3, otherwise
+    2.
+
+    If "C" is present in `axes`, then you need to set `n_channels_in` to the number of
+    channels. Likewise, if you set the number of channels, then "C" must be present in
+    `axes`.
+
+    To set the number of output channels, use the `n_channels_out` parameter. If it is
+    not specified, it will be assumed to be equal to `n_channels_in`.
+
+    By default, all channels are trained together. To train all channels independently,
+    set `independent_channels` to True.
+
+    By setting `use_augmentations` to False, the only transformation applied will be
+    normalization.
+
+    Parameters
+    ----------
+    experiment_name : str
+        Name of the experiment.
+    data_type : Literal["array", "tiff", "custom"]
+        Type of the data.
+    axes : str
+        Axes of the data (e.g. SYX).
+    patch_size : List[int]
+        Size of the patches along the spatial dimensions (e.g. [64, 64]).
+    batch_size : int
+        Batch size.
+    num_epochs : int
+        Number of epochs.
+    use_augmentations : bool, optional
+        Whether to use augmentations, by default True.
+    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_kwargs : dict, optional
+        UNetModel parameters, by default {}.
+
+    Returns
+    -------
+    Configuration
+        Configuration for training CARE.
+    """
+    if n_channels_out == -1:
+        n_channels_out = n_channels_in
+
+    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,
+        use_augmentations=use_augmentations,
+        independent_channels=independent_channels,
+        loss=loss,
+        n_channels_in=n_channels_in,
+        n_channels_out=n_channels_out,
+        logger=logger,
+        model_kwargs=model_kwargs,
+    )
+
+
+def create_n2n_configuration(
+    experiment_name: str,
+    data_type: Literal["array", "tiff", "custom"],
+    axes: str,
+    patch_size: List[int],
+    batch_size: int,
+    num_epochs: int,
+    use_augmentations: bool = True,
+    independent_channels: bool = False,
+    loss: Literal["mae", "mse"] = "mae",
+    n_channels: int = 1,
+    logger: Literal["wandb", "tensorboard", "none"] = "none",
+    model_kwargs: Optional[dict] = None,
+) -> Configuration:
+    """
+    Create a configuration for training Noise2Noise.
+
+    If "Z" is present in `axes`, then `path_size` must be a list of length 3, otherwise
+    2.
+
+    If "C" is present in `axes`, then you need to set `n_channels` to the number of
+    channels. Likewise, if you set the number of channels, then "C" must be present in
+    `axes`.
+
+    By default, all channels are trained together. To train all channels independently,
+    set `independent_channels` to True.
+
+    By setting `use_augmentations` to False, the only transformation applied will be
+    normalization.
+
+    Parameters
+    ----------
+    experiment_name : str
+        Name of the experiment.
+    data_type : Literal["array", "tiff", "custom"]
+        Type of the data.
+    axes : str
+        Axes of the data (e.g. SYX).
+    patch_size : List[int]
+        Size of the patches along the spatial dimensions (e.g. [64, 64]).
+    batch_size : int
+        Batch size.
+    num_epochs : int
+        Number of epochs.
+    use_augmentations : bool, optional
+        Whether to use augmentations, by default True.
+    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 : int, optional
+        Number of channels (in and out), by default 1.
+    logger : Literal["wandb", "tensorboard", "none"], optional
+        Logger to use, by default "none".
+    model_kwargs : dict, optional
+        UNetModel parameters, by default {}.
+
+    Returns
+    -------
+    Configuration
+        Configuration for training Noise2Noise.
+    """
+    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,
+        use_augmentations=use_augmentations,
+        independent_channels=independent_channels,
+        loss=loss,
+        n_channels_in=n_channels,
+        n_channels_out=n_channels,
+        logger=logger,
+        model_kwargs=model_kwargs,
+    )
+
+
 def create_n2v_configuration(
     experiment_name: str,
     data_type: Literal["array", "tiff", "custom"],
@@ -160,8 +328,9 @@ def create_n2v_configuration(
     batch_size: int,
     num_epochs: int,
     use_augmentations: bool = True,
+    independent_channels: bool = True,
     use_n2v2: bool = False,
-    n_channels: int = -1,
+    n_channels: int = 1,
     roi_size: int = 11,
     masked_pixel_percentage: float = 0.2,
     struct_n2v_axis: Literal["horizontal", "vertical", "none"] = "none",
@@ -170,7 +339,7 @@ def create_n2v_configuration(
     model_kwargs: Optional[dict] = None,
 ) -> Configuration:
     """
-    Create a configuration for training N2V.
+    Create a configuration for training Noise2Void.
 
     N2V uses a UNet model to denoise images in a self-supervised manner. To use its
     variants structN2V and N2V2, set the `struct_n2v_axis` and `struct_n2v_span`
@@ -181,11 +350,14 @@ def create_n2v_configuration(
     or horizontal correlations are present in the noise; it applies an additional mask
     to the manipulated pixel neighbors.
 
+    If "Z" is present in `axes`, then `path_size` must be a list of length 3, otherwise
+    2.
+
     If "C" is present in `axes`, then you need to set `n_channels` to the number of
     channels.
 
-    If "Z" is present in `axes`, then `path_size` must be a list of length 3, otherwise
-    2.
+    By default, all channels are trained independently. To train all channels together,
+    set `independent_channels` to False.
 
     By setting `use_augmentations` to False, the only transformations applied will be
     normalization and N2V manipulation.
@@ -217,10 +389,12 @@ def create_n2v_configuration(
         Number of epochs.
     use_augmentations : bool, optional
         Whether to use augmentations, by default True.
+    independent_channels : bool, optional
+        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.
+        Number of channels (in and out), by default 1.
     roi_size : int, optional
         N2V pixel manipulation area, by default 11.
     masked_pixel_percentage : float, optional
@@ -275,8 +449,20 @@ def create_n2v_configuration(
     ...     struct_n2v_span=7
     ... )
 
-    If you are training multiple channels together, then you need to specify the number
-    of channels:
+    If you are training multiple channels independently, then you need to specify the
+    number of channels:
+    >>> config = create_n2v_configuration(
+    ...     experiment_name="n2v_experiment",
+    ...     data_type="array",
+    ...     axes="YXC",
+    ...     patch_size=[64, 64],
+    ...     batch_size=32,
+    ...     num_epochs=100,
+    ...     n_channels=3
+    ... )
+
+    If instead you want to train multiple channels together, you need to turn off the
+    `independent_channels` parameter:
     >>> config = create_n2v_configuration(
     ...     experiment_name="n2v_experiment",
     ...     data_type="array",
@@ -284,6 +470,7 @@ def create_n2v_configuration(
     ...     patch_size=[64, 64],
     ...     batch_size=32,
     ...     num_epochs=100,
+    ...     independent_channels=False,
     ...     n_channels=3
     ... )
 
@@ -300,18 +487,16 @@ def create_n2v_configuration(
     ... )
     """
     # if there are channels, we need to specify their number
-    if "C" in axes and n_channels == -1:
+    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:
+    elif "C" not in axes 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)."
         )
-    elif n_channels == -1:
-        n_channels = 1
 
     # model
     if model_kwargs is None:
@@ -320,6 +505,7 @@ def create_n2v_configuration(
     model_kwargs["conv_dims"] = 3 if "Z" in axes else 2
     model_kwargs["in_channels"] = n_channels
     model_kwargs["num_classes"] = n_channels
+    model_kwargs["independent_channels"] = independent_channels
 
     unet_model = UNetModel(
         architecture=SupportedArchitecture.UNET.value,
@@ -327,7 +513,7 @@ def create_n2v_configuration(
     )
 
     # algorithm model
-    algorithm = AlgorithmModel(
+    algorithm = AlgorithmConfig(
         algorithm=SupportedAlgorithm.N2V.value,
         loss=SupportedLoss.N2V.value,
         model=unet_model,
@@ -356,9 +542,11 @@ def create_n2v_configuration(
     # n2v2 and structn2v
     nv2_transform = {
         "name": SupportedTransform.N2V_MANIPULATE.value,
-        "strategy": SupportedPixelManipulation.MEDIAN.value
-        if use_n2v2
-        else SupportedPixelManipulation.UNIFORM.value,
+        "strategy": (
+            SupportedPixelManipulation.MEDIAN.value
+            if use_n2v2
+            else SupportedPixelManipulation.UNIFORM.value
+        ),
         "roi_size": roi_size,
         "masked_pixel_percentage": masked_pixel_percentage,
         "struct_mask_axis": struct_n2v_axis,
@@ -367,7 +555,7 @@ def create_n2v_configuration(
     transforms.append(nv2_transform)
 
     # data model
-    data = DataModel(
+    data = DataConfig(
         data_type=data_type,
         axes=axes,
         patch_size=patch_size,
@@ -376,7 +564,7 @@ def create_n2v_configuration(
     )
 
     # training model
-    training = TrainingModel(
+    training = TrainingConfig(
         num_epochs=num_epochs,
         batch_size=batch_size,
         logger=None if logger == "none" else logger,
@@ -393,17 +581,16 @@ def create_n2v_configuration(
     return configuration
 
 
-# TODO add tests
 def create_inference_configuration(
-    training_configuration: Configuration,
+    configuration: Configuration,
     tile_size: Optional[Tuple[int, ...]] = None,
     tile_overlap: Optional[Tuple[int, ...]] = None,
     data_type: Optional[Literal["array", "tiff", "custom"]] = None,
     axes: Optional[str] = None,
-    transforms: Optional[Union[List[Dict[str, Any]], Compose]] = None,
+    transforms: Optional[Union[List[Dict[str, Any]]]] = None,
     tta_transforms: bool = True,
     batch_size: Optional[int] = 1,
-) -> InferenceModel:
+) -> InferenceConfig:
     """
     Create a configuration for inference with N2V.
 
@@ -412,8 +599,8 @@ def create_inference_configuration(
 
     Parameters
     ----------
-    training_configuration : Configuration
-        Configuration used for training.
+    configuration : Configuration
+        Global configuration.
     tile_size : Tuple[int, ...], optional
         Size of the tiles.
     tile_overlap : Tuple[int, ...], optional
@@ -422,7 +609,7 @@ def create_inference_configuration(
         Type of the data, by default "tiff".
     axes : str, optional
         Axes of the data, by default "YX".
-    transforms : List[Dict[str, Any]] or Compose, optional
+    transforms : List[Dict[str, Any]], optional
         Transformations to apply to the data, by default None.
     tta_transforms : bool, optional
         Whether to apply test-time augmentations, by default True.
@@ -432,14 +619,12 @@ def create_inference_configuration(
     Returns
     -------
     InferenceConfiguration
-        Configuration for inference with N2V.
+        Configuration used to configure CAREamicsPredictData.
     """
-    if (
-        training_configuration.data_config.mean is None
-        or training_configuration.data_config.std is None
-    ):
-        raise ValueError("Mean and std must be provided in the training configuration.")
+    if configuration.data_config.mean is None or configuration.data_config.std is None:
+        raise ValueError("Mean and std must be provided in the configuration.")
 
+    # minimum transform
     if transforms is None:
         transforms = [
             {
@@ -447,13 +632,35 @@ def create_inference_configuration(
             },
         ]
 
-    return InferenceModel(
-        data_type=data_type or training_configuration.data_config.data_type,
+    # tile size for UNets
+    if tile_size is not None:
+        model = configuration.algorithm_config.model
+
+        if model.architecture == SupportedArchitecture.UNET.value:
+            # tile size must be equal to k*2^n, where n is the number of pooling layers
+            # (equal to the depth) and k is an integer
+            depth = model.depth
+            tile_increment = 2**depth
+
+            for i, t in enumerate(tile_size):
+                if t % tile_increment != 0:
+                    raise ValueError(
+                        f"Tile size must be divisible by {tile_increment} along all "
+                        f"axes (got {t} for axis {i}). If your image size is smaller "
+                        f"along one axis (e.g. Z), consider padding the image."
+                    )
+
+        # tile overlaps must be specified
+        if tile_overlap is None:
+            raise ValueError("Tile overlap must be specified.")
+
+    return InferenceConfig(
+        data_type=data_type or configuration.data_config.data_type,
         tile_size=tile_size,
         tile_overlap=tile_overlap,
-        axes=axes or training_configuration.data_config.axes,
-        mean=training_configuration.data_config.mean,
-        std=training_configuration.data_config.std,
+        axes=axes or configuration.data_config.axes,
+        mean=configuration.data_config.mean,
+        std=configuration.data_config.std,
         transforms=transforms,
         tta_transforms=tta_transforms,
         batch_size=batch_size,

careamics/config/configuration_model.py

@@ -1,4 +1,5 @@
 """Pydantic CAREamics configuration."""
+
 from __future__ import annotations
 
 import re
@@ -11,8 +12,8 @@ from bioimageio.spec.generic.v0_3 import CiteEntry
 from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 from typing_extensions import Self
 
-from .algorithm_model import AlgorithmModel
-from .data_model import DataModel
+from .algorithm_model import AlgorithmConfig
+from .data_model import DataConfig
 from .references import (
     CARE,
     CUSTOM,
@@ -34,7 +35,7 @@ from .references import (
     StructN2VRef,
 )
 from .support import SupportedAlgorithm, SupportedPixelManipulation, SupportedTransform
-from .training_model import TrainingModel
+from .training_model import TrainingConfig
 from .transformations.n2v_manipulate_model import (
     N2VManipulateModel,
 )
@@ -156,9 +157,10 @@ class Configuration(BaseModel):
     )
 
     # Sub-configurations
-    algorithm_config: AlgorithmModel
-    data_config: DataModel
-    training_config: TrainingModel
+    algorithm_config: AlgorithmConfig
+
+    data_config: DataConfig
+    training_config: TrainingConfig
 
     @field_validator("experiment_name")
     @classmethod
@@ -237,25 +239,22 @@ class Configuration(BaseModel):
             Validated configuration.
         """
         if self.algorithm_config.algorithm == SupportedAlgorithm.N2V:
-            # if we have a list of transform (as opposed to Compose)
-            if self.data_config.has_transform_list():
-                # missing N2V_MANIPULATE
-                if not self.data_config.has_n2v_manipulate():
-                    self.data_config.transforms.append(
-                        N2VManipulateModel(
-                            name=SupportedTransform.N2V_MANIPULATE.value,
-                        )
+            # missing N2V_MANIPULATE
+            if not self.data_config.has_n2v_manipulate():
+                self.data_config.transforms.append(
+                    N2VManipulateModel(
+                        name=SupportedTransform.N2V_MANIPULATE.value,
                     )
+                )
 
-                median = SupportedPixelManipulation.MEDIAN.value
-                uniform = SupportedPixelManipulation.UNIFORM.value
-                strategy = median if self.algorithm_config.model.n2v2 else uniform
-                self.data_config.set_N2V2_strategy(strategy)
+            median = SupportedPixelManipulation.MEDIAN.value
+            uniform = SupportedPixelManipulation.UNIFORM.value
+            strategy = median if self.algorithm_config.model.n2v2 else uniform
+            self.data_config.set_N2V2_strategy(strategy)
         else:
-            # if we have a list of transform, remove N2V manipulate if present
-            if self.data_config.has_transform_list():
-                if self.data_config.has_n2v_manipulate():
-                    self.data_config.remove_n2v_manipulate()
+            # remove N2V manipulate if present
+            if self.data_config.has_n2v_manipulate():
+                self.data_config.remove_n2v_manipulate()
 
         return self
 
@@ -591,6 +590,6 @@ def save_configuration(config: Configuration, path: Union[str, Path]) -> 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)
+        yaml.dump(config.model_dump(), f, default_flow_style=False, sort_keys=False)
 
     return config_path