careamics 0.1.0rc6__py3-none-any.whl → 0.1.0rc8__py3-none-any.whl

careamics/config/__init__.py

@@ -14,9 +14,7 @@ __all__ = [
     "create_care_configuration",
     "register_model",
     "CustomModel",
-    "create_inference_configuration",
     "clear_custom_models",
-    "ConfigurationInformation",
 ]
 
 from .algorithm_model import AlgorithmConfig
@@ -24,7 +22,6 @@ from .architectures import CustomModel, clear_custom_models, register_model
 from .callback_model import CheckpointModel
 from .configuration_factory import (
     create_care_configuration,
-    create_inference_configuration,
     create_n2n_configuration,
     create_n2v_configuration,
 )

careamics/config/algorithm_model.py

@@ -93,12 +93,20 @@ class AlgorithmConfig(BaseModel):
 
     # Mandatory fields
     algorithm: Literal["n2v", "care", "n2n", "custom"]  # defined in SupportedAlgorithm
+    """Name of the algorithm, as defined in SupportedAlgorithm."""
+
     loss: Literal["n2v", "mae", "mse"]
+    """Loss function to use, as defined in SupportedLoss."""
+
     model: Union[UNetModel, VAEModel, CustomModel] = Field(discriminator="architecture")
+    """Model architecture to use, defined in SupportedArchitecture."""
 
     # Optional fields
     optimizer: OptimizerModel = OptimizerModel()
+    """Optimizer to use, defined in SupportedOptimizer."""
+
     lr_scheduler: LrSchedulerModel = LrSchedulerModel()
+    """Learning rate scheduler to use, defined in SupportedScheduler."""
 
     @model_validator(mode="after")
     def algorithm_cross_validation(self: Self) -> Self:
@@ -134,21 +142,6 @@ class AlgorithmConfig(BaseModel):
                     "sure that `in_channels` and `num_classes` are the same."
                 )
 
-        # N2N
-        if self.algorithm == "n2n":
-            # n2n is only compatible with the UNet model
-            if not isinstance(self.model, UNetModel):
-                raise ValueError(
-                    f"Model for algorithm {self.algorithm} must be a `UNetModel`."
-                )
-
-            # n2n requires the number of input and output channels to be the same
-            if self.model.in_channels != self.model.num_classes:
-                raise ValueError(
-                    "N2N requires the same number of input and output channels. Make "
-                    "sure that `in_channels` and `num_classes` are the same."
-                )
-
         if self.algorithm == "care" or self.algorithm == "n2n":
             if self.loss == "n2v":
                 raise ValueError("Supervised algorithms do not support loss `n2v`.")

careamics/config/architectures/architecture_model.py

@@ -13,6 +13,7 @@ class ArchitectureModel(BaseModel):
     """
 
     architecture: str
+    """Name of the architecture."""
 
     def model_dump(self, **kwargs: Any) -> Dict[str, Any]:
         """

careamics/config/architectures/custom_model.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from pprint import pformat
-from typing import Any, Dict, Literal
+from typing import Any, Literal
 
 from pydantic import ConfigDict, field_validator, model_validator
 from torch.nn import Module
@@ -72,9 +72,11 @@ class CustomModel(ArchitectureModel):
 
     # discriminator used for choosing the pydantic model in Model
     architecture: Literal["Custom"]
+    """Name of the architecture."""
 
     # name of the custom model
     name: str
+    """Name of the custom model."""
 
     @field_validator("name")
     @classmethod
@@ -136,7 +138,7 @@ class CustomModel(ArchitectureModel):
         """
         return pformat(self.model_dump())
 
-    def model_dump(self, **kwargs: Any) -> Dict[str, Any]:
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
         """Dump the model configuration.
 
         Parameters
@@ -146,7 +148,7 @@ class CustomModel(ArchitectureModel):
 
         Returns
         -------
-        Dict[str, Any]
+        dict[str, Any]
             Model configuration.
         """
         model_dict = super().model_dump()

careamics/config/architectures/unet_model.py

@@ -29,19 +29,38 @@ class UNetModel(ArchitectureModel):
 
     # discriminator used for choosing the pydantic model in Model
     architecture: Literal["UNet"]
+    """Name of the architecture."""
 
     # parameters
     # validate_defaults allow ignoring default values in the dump if they were not set
     conv_dims: Literal[2, 3] = Field(default=2, validate_default=True)
+    """Dimensions (2D or 3D) of the convolutional layers."""
+
     num_classes: int = Field(default=1, ge=1, validate_default=True)
+    """Number of classes or channels in the model output."""
+
     in_channels: int = Field(default=1, ge=1, validate_default=True)
+    """Number of channels in the input to the model."""
+
     depth: int = Field(default=2, ge=1, le=10, validate_default=True)
+    """Number of levels in the UNet."""
+
     num_channels_init: int = Field(default=32, ge=8, le=1024, validate_default=True)
+    """Number of convolutional filters in the first layer of the UNet."""
+
     final_activation: Literal[
         "None", "Sigmoid", "Softmax", "Tanh", "ReLU", "LeakyReLU"
     ] = Field(default="None", validate_default=True)
+    """Final activation function."""
+
     n2v2: bool = Field(default=False, validate_default=True)
+    """Whether to use N2V2 architecture modifications, with blur pool layers and fewer
+    skip connections.
+    """
+
     independent_channels: bool = Field(default=True, validate_default=True)
+    """Whether information is processed independently in each channel, used to train
+    channels independently."""
 
     @field_validator("num_channels_init")
     @classmethod

careamics/config/architectures/vae_model.py

@@ -17,6 +17,7 @@ class VAEModel(ArchitectureModel):
     )
 
     architecture: Literal["VAE"]
+    """Name of the architecture."""
 
     def set_3D(self, is_3D: bool) -> None:
         """

careamics/config/callback_model.py

@@ -13,69 +13,111 @@ from pydantic import (
 
 
 class CheckpointModel(BaseModel):
-    """Checkpoint saving callback Pydantic model."""
+    """Checkpoint saving callback Pydantic model.
+
+    The parameters corresponds to those of
+    `pytorch_lightning.callbacks.ModelCheckpoint`.
+
+    See:
+    https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html#modelcheckpoint
+    """
 
     model_config = ConfigDict(
         validate_assignment=True,
     )
 
     monitor: Literal["val_loss"] = Field(default="val_loss", validate_default=True)
+    """Quantity to monitor."""
+
     verbose: bool = Field(default=False, validate_default=True)
+    """Verbosity mode."""
+
     save_weights_only: bool = Field(default=False, validate_default=True)
+    """When `True`, only the model's weights will be saved (model.save_weights)."""
+
+    save_last: Optional[Literal[True, False, "link"]] = Field(
+        default=True, validate_default=True
+    )
+    """When `True`, saves a last.ckpt copy whenever a checkpoint file gets saved."""
+
+    save_top_k: int = Field(default=3, ge=1, le=10, validate_default=True)
+    """If `save_top_k == kz, the best k models according to the quantity monitored
+    will be saved. If `save_top_k == 0`, no models are saved. if `save_top_k == -1`,
+    all models are saved."""
+
     mode: Literal["min", "max"] = Field(default="min", validate_default=True)
+    """One of {min, max}. If `save_top_k != 0`, the decision to overwrite the current
+    save file is made based on either the maximization or the minimization of the
+    monitored quantity. For 'val_acc', this should be 'max', for 'val_loss' this should
+    be 'min', etc.
+    """
+
     auto_insert_metric_name: bool = Field(default=False, validate_default=True)
+    """When `True`, the checkpoints filenames will contain the metric name."""
+
     every_n_train_steps: Optional[int] = Field(
         default=None, ge=1, le=10, validate_default=True
     )
+    """Number of training steps between checkpoints."""
+
     train_time_interval: Optional[timedelta] = Field(
         default=None, validate_default=True
     )
+    """Checkpoints are monitored at the specified time interval."""
+
     every_n_epochs: Optional[int] = Field(
         default=None, ge=1, le=10, validate_default=True
     )
-    save_last: Optional[Literal[True, False, "link"]] = Field(
-        default=True, validate_default=True
-    )
-    save_top_k: int = Field(default=3, ge=1, le=10, validate_default=True)
+    """Number of epochs between checkpoints."""
 
 
 class EarlyStoppingModel(BaseModel):
-    """Early stopping callback Pydantic model."""
+    """Early stopping callback Pydantic model.
+
+    The parameters corresponds to those of
+    `pytorch_lightning.callbacks.ModelCheckpoint`.
+
+    See:
+    https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.EarlyStopping.html#lightning.pytorch.callbacks.EarlyStopping
+    """
 
     model_config = ConfigDict(
         validate_assignment=True,
     )
 
     monitor: Literal["val_loss"] = Field(default="val_loss", validate_default=True)
+    """Quantity to monitor."""
+
+    min_delta: float = Field(default=0.0, ge=0.0, le=1.0, validate_default=True)
+    """Minimum change in the monitored quantity to qualify as an improvement, i.e. an
+    absolute change of less than or equal to min_delta, will count as no improvement."""
+
     patience: int = Field(default=3, ge=1, le=10, validate_default=True)
+    """Number of checks with no improvement after which training will be stopped."""
+
+    verbose: bool = Field(default=False, validate_default=True)
+    """Verbosity mode."""
+
     mode: Literal["min", "max", "auto"] = Field(default="min", validate_default=True)
-    min_delta: float = Field(default=0.0, ge=0.0, le=1.0, validate_default=True)
+    """One of {min, max, auto}."""
+
     check_finite: bool = Field(default=True, validate_default=True)
-    stop_on_nan: bool = Field(default=True, validate_default=True)
-    verbose: bool = Field(default=False, validate_default=True)
-    restore_best_weights: bool = Field(default=True, validate_default=True)
-    auto_lr_find: bool = Field(default=False, validate_default=True)
-    auto_lr_find_patience: int = Field(default=3, ge=1, le=10, validate_default=True)
-    auto_lr_find_mode: Literal["min", "max", "auto"] = Field(
-        default="min", validate_default=True
-    )
-    auto_lr_find_direction: Literal["forward", "backward"] = Field(
-        default="backward", validate_default=True
-    )
-    auto_lr_find_max_lr: float = Field(
-        default=10.0, ge=0.0, le=1e6, validate_default=True
-    )
-    auto_lr_find_min_lr: float = Field(
-        default=1e-8, ge=0.0, le=1e6, validate_default=True
-    )
-    auto_lr_find_num_training: int = Field(
-        default=100, ge=1, le=1e6, validate_default=True
-    )
-    auto_lr_find_divergence_threshold: float = Field(
-        default=5.0, ge=0.0, le=1e6, validate_default=True
-    )
-    auto_lr_find_accumulate_grad_batches: int = Field(
-        default=1, ge=1, le=1e6, validate_default=True
+    """When `True`, stops training when the monitored quantity becomes `NaN` or
+    `inf`."""
+
+    stopping_threshold: Optional[float] = Field(default=None, validate_default=True)
+    """Stop training immediately once the monitored quantity reaches this threshold."""
+
+    divergence_threshold: Optional[float] = Field(default=None, validate_default=True)
+    """Stop training as soon as the monitored quantity becomes worse than this
+    threshold."""
+
+    check_on_train_epoch_end: Optional[bool] = Field(
+        default=False, validate_default=True
     )
-    auto_lr_find_stop_divergence: bool = Field(default=True, validate_default=True)
-    auto_lr_find_step_scale: float = Field(default=0.1, ge=0.0, le=10)
+    """Whether to run early stopping at the end of the training epoch. If this is
+    `False`, then the check runs at the end of the validation."""
+
+    log_rank_zero_only: bool = Field(default=False, validate_default=True)
+    """When set `True`, logs the status of the early stopping callback only for rank 0
+    process."""

careamics/config/configuration_factory.py

@@ -1,12 +1,11 @@
 """Convenience functions to create configurations for training and inference."""
 
-from typing import Any, Dict, List, Literal, Optional, Tuple
+from typing import Any, Dict, List, Literal, Optional
 
 from .algorithm_model import AlgorithmConfig
 from .architectures import UNetModel
 from .configuration_model import Configuration
 from .data_model import DataConfig
-from .inference_model import InferenceConfig
 from .support import (
     SupportedAlgorithm,
     SupportedArchitecture,
@@ -107,9 +106,6 @@ def _create_supervised_configuration(
     # augmentations
     if use_augmentations:
         transforms: List[Dict[str, Any]] = [
-            {
-                "name": SupportedTransform.NORMALIZE.value,
-            },
             {
                 "name": SupportedTransform.XY_FLIP.value,
             },
@@ -118,11 +114,7 @@ def _create_supervised_configuration(
             },
         ]
     else:
-        transforms = [
-            {
-                "name": SupportedTransform.NORMALIZE.value,
-            },
-        ]
+        transforms = []
 
     # data model
     data = DataConfig(
@@ -250,7 +242,8 @@ def create_n2n_configuration(
     use_augmentations: bool = True,
     independent_channels: bool = False,
     loss: Literal["mae", "mse"] = "mae",
-    n_channels: int = 1,
+    n_channels_in: int = 1,
+    n_channels_out: int = -1,
     logger: Literal["wandb", "tensorboard", "none"] = "none",
     model_kwargs: Optional[dict] = None,
 ) -> Configuration:
@@ -260,10 +253,13 @@ def create_n2n_configuration(
     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
+    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.
 
@@ -290,8 +286,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 : int, optional
-        Number of channels (in and out), by default 1.
+    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
@@ -302,6 +300,9 @@ def create_n2n_configuration(
     Configuration
         Configuration for training Noise2Noise.
     """
+    if n_channels_out == -1:
+        n_channels_out = n_channels_in
+
     return _create_supervised_configuration(
         algorithm="n2n",
         experiment_name=experiment_name,
@@ -313,8 +314,8 @@ def create_n2n_configuration(
         use_augmentations=use_augmentations,
         independent_channels=independent_channels,
         loss=loss,
-        n_channels_in=n_channels,
-        n_channels_out=n_channels,
+        n_channels_in=n_channels_in,
+        n_channels_out=n_channels_out,
         logger=logger,
         model_kwargs=model_kwargs,
     )
@@ -522,9 +523,6 @@ def create_n2v_configuration(
     # augmentations
     if use_augmentations:
         transforms: List[Dict[str, Any]] = [
-            {
-                "name": SupportedTransform.NORMALIZE.value,
-            },
             {
                 "name": SupportedTransform.XY_FLIP.value,
             },
@@ -533,11 +531,7 @@ def create_n2v_configuration(
             },
         ]
     else:
-        transforms = [
-            {
-                "name": SupportedTransform.NORMALIZE.value,
-            },
-        ]
+        transforms = []
 
     # n2v2 and structn2v
     nv2_transform = {
@@ -579,77 +573,3 @@ def create_n2v_configuration(
     )
 
     return configuration
-
-
-def create_inference_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,
-    tta_transforms: bool = True,
-    batch_size: Optional[int] = 1,
-) -> InferenceConfig:
-    """
-    Create a configuration for inference with N2V.
-
-    If not provided, `data_type` and `axes` are taken from the training
-    configuration.
-
-    Parameters
-    ----------
-    configuration : Configuration
-        Global configuration.
-    tile_size : Tuple[int, ...], optional
-        Size of the tiles.
-    tile_overlap : Tuple[int, ...], optional
-        Overlap of the tiles.
-    data_type : str, optional
-        Type of the data, by default "tiff".
-    axes : str, optional
-        Axes of the data, by default "YX".
-    tta_transforms : bool, optional
-        Whether to apply test-time augmentations, by default True.
-    batch_size : int, optional
-        Batch size, by default 1.
-
-    Returns
-    -------
-    InferenceConfiguration
-        Configuration used to configure CAREamicsPredictData.
-    """
-    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.")
-
-    # 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 configuration.data_config.axes,
-        mean=configuration.data_config.mean,
-        std=configuration.data_config.std,
-        tta_transforms=tta_transforms,
-        batch_size=batch_size,
-    )

careamics/config/configuration_model.py

@@ -5,11 +5,11 @@ from __future__ import annotations
 import re
 from pathlib import Path
 from pprint import pformat
-from typing import Dict, List, Literal, Union
+from typing import Literal, Union
 
 import yaml
 from bioimageio.spec.generic.v0_3 import CiteEntry
-from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
+from pydantic import BaseModel, ConfigDict, field_validator, model_validator
 from typing_extensions import Self
 
 from .algorithm_model import AlgorithmConfig
@@ -147,20 +147,25 @@ class Configuration(BaseModel):
     )
 
     # version
-    version: Literal["0.1.0"] = Field(
-        default="0.1.0", description="Version of the CAREamics configuration."
-    )
+    version: Literal["0.1.0"] = "0.1.0"
+    """CAREamics configuration version."""
 
     # required parameters
-    experiment_name: str = Field(
-        ..., description="Name of the experiment, used to name logs and checkpoints."
-    )
+    experiment_name: str
+    """Name of the experiment, used to name logs and checkpoints."""
 
     # Sub-configurations
     algorithm_config: AlgorithmConfig
+    """Algorithm configuration, holding all parameters required to configure the
+    model."""
 
     data_config: DataConfig
+    """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
@@ -269,7 +274,7 @@ class Configuration(BaseModel):
         """
         return pformat(self.model_dump())
 
-    def set_3D(self, is_3D: bool, axes: str, patch_size: List[int]) -> None:
+    def set_3D(self, is_3D: bool, axes: str, patch_size: list[int]) -> None:
         """
         Set 3D flag and axes.
 
@@ -279,7 +284,7 @@ class Configuration(BaseModel):
             Whether the algorithm is 3D or not.
         axes : str
             Axes of the data.
-        patch_size : List[int]
+        patch_size : list[int]
             Patch size.
         """
         # set the flag and axes (this will not trigger validation at the config level)
@@ -389,7 +394,7 @@ class Configuration(BaseModel):
 
         return ""
 
-    def get_algorithm_citations(self) -> List[CiteEntry]:
+    def get_algorithm_citations(self) -> list[CiteEntry]:
         """
         Return a list of citation entries of the current algorithm.
 
@@ -455,13 +460,13 @@ class Configuration(BaseModel):
 
         return ""
 
-    def get_algorithm_keywords(self) -> List[str]:
+    def get_algorithm_keywords(self) -> list[str]:
         """
         Get algorithm keywords.
 
         Returns
         -------
-        List[str]
+        list[str]
             List of keywords.
         """
         if self.algorithm_config.algorithm == SupportedAlgorithm.N2V:
@@ -491,8 +496,8 @@ class Configuration(BaseModel):
         self,
         exclude_defaults: bool = False,
         exclude_none: bool = True,
-        **kwargs: Dict,
-    ) -> Dict:
+        **kwargs: dict,
+    ) -> dict:
         """
         Override model_dump method in order to set default values.
 
@@ -503,7 +508,7 @@ class Configuration(BaseModel):
             True.
         exclude_none : bool, optional
             Whether to exclude fields with None values or not, by default True.
-        **kwargs : Dict
+        **kwargs : dict
             Keyword arguments.
 
         Returns
@@ -524,7 +529,7 @@ def load_configuration(path: Union[str, Path]) -> Configuration:
 
     Parameters
     ----------
-    path : Union[str, Path]
+    path : str or Path
         Path to the configuration.
 
     Returns
@@ -556,7 +561,7 @@ def save_configuration(config: Configuration, path: Union[str, Path]) -> Path:
     ----------
     config : Configuration
         Configuration to save.
-    path : Union[str, Path]
+    path : str or Path
         Path to a existing folder in which to save the configuration or to an existing
         configuration file.