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

careamics/config/algorithms/care_algorithm_model.py

@@ -1,10 +1,14 @@
 """CARE algorithm configuration."""
 
-from typing import Literal
+from typing import Annotated, Literal
 
-from pydantic import field_validator
+from pydantic import AfterValidator
 
 from careamics.config.architectures import UNetModel
+from careamics.config.validators import (
+    model_without_final_activation,
+    model_without_n2v2,
+)
 
 from .unet_algorithm_model import UNetBasedAlgorithm
 
@@ -26,25 +30,9 @@ class CAREAlgorithm(UNetBasedAlgorithm):
     loss: Literal["mae", "mse"] = "mae"
     """CARE-compatible loss function."""
 
-    @classmethod
-    @field_validator("model")
-    def model_without_n2v2(cls, value: UNetModel) -> UNetModel:
-        """Validate that the model does not have the n2v2 attribute.
-
-        Parameters
-        ----------
-        value : UNetModel
-            Model to validate.
-
-        Returns
-        -------
-        UNetModel
-            The validated model.
-        """
-        if value.n2v2:
-            raise ValueError(
-                "The N2N algorithm does not support the `n2v2` attribute. "
-                "Set it to `False`."
-            )
-
-        return value
+    model: Annotated[
+        UNetModel,
+        AfterValidator(model_without_n2v2),
+        AfterValidator(model_without_final_activation),
+    ]
+    """UNet without a final activation function and without the `n2v2` modifications."""

careamics/config/algorithms/n2n_algorithm_model.py

@@ -1,16 +1,20 @@
 """N2N Algorithm configuration."""
 
-from typing import Literal
+from typing import Annotated, Literal
 
-from pydantic import field_validator
+from pydantic import AfterValidator
 
 from careamics.config.architectures import UNetModel
+from careamics.config.validators import (
+    model_without_final_activation,
+    model_without_n2v2,
+)
 
 from .unet_algorithm_model import UNetBasedAlgorithm
 
 
 class N2NAlgorithm(UNetBasedAlgorithm):
-    """N2N Algorithm configuration."""
+    """Noise2Noise Algorithm configuration."""
 
     algorithm: Literal["n2n"] = "n2n"
     """N2N Algorithm name."""
@@ -18,25 +22,9 @@ class N2NAlgorithm(UNetBasedAlgorithm):
     loss: Literal["mae", "mse"] = "mae"
     """N2N-compatible loss function."""
 
-    @classmethod
-    @field_validator("model")
-    def model_without_n2v2(cls, value: UNetModel) -> UNetModel:
-        """Validate that the model does not have the n2v2 attribute.
-
-        Parameters
-        ----------
-        value : UNetModel
-            Model to validate.
-
-        Returns
-        -------
-        UNetModel
-            The validated model.
-        """
-        if value.n2v2:
-            raise ValueError(
-                "The N2N algorithm does not support the `n2v2` attribute. "
-                "Set it to `False`."
-            )
-
-        return value
+    model: Annotated[
+        UNetModel,
+        AfterValidator(model_without_n2v2),
+        AfterValidator(model_without_final_activation),
+    ]
+    """UNet without a final activation function and without the `n2v2` modifications."""

careamics/config/algorithms/n2v_algorithm_model.py

@@ -1,9 +1,14 @@
 """"N2V Algorithm configuration."""
 
-from typing import Literal
+from typing import Annotated, Literal
 
-from pydantic import model_validator
-from typing_extensions import Self
+from pydantic import AfterValidator
+
+from careamics.config.architectures import UNetModel
+from careamics.config.validators import (
+    model_matching_in_out_channels,
+    model_without_final_activation,
+)
 
 from .unet_algorithm_model import UNetBasedAlgorithm
 
@@ -17,19 +22,8 @@ class N2VAlgorithm(UNetBasedAlgorithm):
     loss: Literal["n2v"] = "n2v"
     """N2V loss function."""
 
-    @model_validator(mode="after")
-    def algorithm_cross_validation(self: Self) -> Self:
-        """Validate the algorithm model for N2V.
-
-        Returns
-        -------
-        Self
-            The validated model.
-        """
-        if self.model.in_channels != self.model.num_classes:
-            raise ValueError(
-                "N2V requires the same number of input and output channels. Make "
-                "sure that `in_channels` and `num_classes` are the same."
-            )
-
-        return self
+    model: Annotated[
+        UNetModel,
+        AfterValidator(model_matching_in_out_channels),
+        AfterValidator(model_without_final_activation),
+    ]

careamics/config/configuration_factories.py

@@ -1,8 +1,8 @@
 """Convenience functions to create configurations for training and inference."""
 
-from typing import Any, Literal, Optional, Union
+from typing import Annotated, Any, Literal, Optional, Union
 
-from pydantic import TypeAdapter
+from pydantic import Discriminator, Tag, TypeAdapter
 
 from careamics.config.algorithms import CAREAlgorithm, N2NAlgorithm, N2VAlgorithm
 from careamics.config.architectures import UNetModel
@@ -12,6 +12,7 @@ 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,
@@ -26,6 +27,24 @@ from careamics.config.transformations import (
 )
 
 
+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]:
@@ -43,7 +62,14 @@ def configuration_factory(
         Configuration for training CAREamics.
     """
     adapter: TypeAdapter = TypeAdapter(
-        Union[N2VConfiguration, N2NConfiguration, CAREConfiguration]
+        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)
 
@@ -198,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.
@@ -236,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
     -------
@@ -268,8 +297,12 @@ def _create_configuration(
         "patch_size": patch_size,
         "batch_size": batch_size,
         "transforms": augmentations,
-        "dataloader_params": dataloader_params,
     }
+    # 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(
@@ -305,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.
@@ -342,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
     -------
@@ -390,7 +426,8 @@ def _create_supervised_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,
     )
 
 
@@ -408,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.
@@ -461,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
     -------
@@ -551,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,
     )
 
 
@@ -569,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.
@@ -622,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
     -------
@@ -712,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,
     )
 
 
@@ -733,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.
@@ -812,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
     -------
@@ -953,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/data/data_model.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 from collections.abc import Sequence
 from pprint import pformat
 from typing import Annotated, Any, Literal, Optional, Union
+from warnings import warn
 
 import numpy as np
 from numpy.typing import NDArray
@@ -100,8 +101,13 @@ class GeneralDataConfig(BaseModel):
     """List of transformations to apply to the data, available transforms are defined
     in SupportedTransform."""
 
-    dataloader_params: Optional[dict] = None
-    """Dictionary of PyTorch dataloader parameters."""
+    train_dataloader_params: dict[str, Any] = Field(
+        default={"shuffle": True}, validate_default=True
+    )
+    """Dictionary of PyTorch training dataloader parameters."""
+
+    val_dataloader_params: dict[str, Any] = Field(default={})
+    """Dictionary of PyTorch validation dataloader parameters."""
 
     @field_validator("patch_size")
     @classmethod
@@ -167,6 +173,45 @@ class GeneralDataConfig(BaseModel):
 
         return axes
 
+    @field_validator("train_dataloader_params")
+    @classmethod
+    def shuffle_train_dataloader(
+        cls, train_dataloader_params: dict[str, Any]
+    ) -> dict[str, Any]:
+        """
+        Validate that "shuffle" is included in the training dataloader params.
+
+        A warning will be raised if `shuffle=False`.
+
+        Parameters
+        ----------
+        train_dataloader_params : dict of {str: Any}
+            The training dataloader parameters.
+
+        Returns
+        -------
+        dict of {str: Any}
+            The validated training dataloader parameters.
+
+        Raises
+        ------
+        ValueError
+            If "shuffle" is not included in the training dataloader params.
+        """
+        if "shuffle" not in train_dataloader_params:
+            raise ValueError(
+                "Value for 'shuffle' was not included in the `train_dataloader_params`."
+            )
+        elif ("shuffle" in train_dataloader_params) and (
+            not train_dataloader_params["shuffle"]
+        ):
+            warn(
+                "Dataloader parameters include `shuffle=False`, this will be passed to "
+                "the training dataloader and may result in bad results.",
+                stacklevel=1,
+            )
+        return train_dataloader_params
+
     @model_validator(mode="after")
     def std_only_with_mean(self: Self) -> Self:
         """

careamics/config/support/supported_algorithms.py

@@ -6,7 +6,11 @@ from careamics.utils import BaseEnum
 
 
 class SupportedAlgorithm(str, BaseEnum):
-    """Algorithms available in CAREamics."""
+    """Algorithms available in CAREamics.
+
+    These definitions are the same as the keyword `name` of the algorithm
+    configurations.
+    """
 
     N2V = "n2v"
     """Noise2Void algorithm, a self-supervised approach based on blind denoising."""

careamics/config/validators/__init__.py

@@ -1,5 +1,16 @@
 """Validator utilities."""
 
-__all__ = ["check_axes_validity", "patch_size_ge_than_8_power_of_2"]
+__all__ = [
+    "check_axes_validity",
+    "model_matching_in_out_channels",
+    "model_without_final_activation",
+    "model_without_n2v2",
+    "patch_size_ge_than_8_power_of_2",
+]
 
+from .model_validators import (
+    model_matching_in_out_channels,
+    model_without_final_activation,
+    model_without_n2v2,
+)
 from .validator_utils import check_axes_validity, patch_size_ge_than_8_power_of_2

careamics/config/validators/model_validators.py

@@ -0,0 +1,84 @@
+"""Architecture model validators."""
+
+from careamics.config.architectures import UNetModel
+
+
+def model_without_n2v2(model: UNetModel) -> UNetModel:
+    """Validate that the Unet model does not have the n2v2 attribute.
+
+    Parameters
+    ----------
+    model : UNetModel
+        Model to validate.
+
+    Returns
+    -------
+    UNetModel
+        The validated model.
+
+    Raises
+    ------
+    ValueError
+        If the model has the `n2v2` attribute set to `True`.
+    """
+    if model.n2v2:
+        raise ValueError(
+            "The algorithm does not support the `n2v2` attribute in the model. "
+            "Set it to `False`."
+        )
+
+    return model
+
+
+def model_without_final_activation(model: UNetModel) -> UNetModel:
+    """Validate that the UNet model does not have the final_activation.
+
+    Parameters
+    ----------
+    model : UNetModel
+        Model to validate.
+
+    Returns
+    -------
+    UNetModel
+        The validated model.
+
+    Raises
+    ------
+    ValueError
+        If the model has the final_activation attribute set.
+    """
+    if model.final_activation != "None":
+        raise ValueError(
+            "The algorithm does not support a `final_activation` in the model. "
+            'Set it to `"None"`.'
+        )
+
+    return model
+
+
+def model_matching_in_out_channels(model: UNetModel) -> UNetModel:
+    """Validate that the UNet model has the same number of channel inputs and outputs.
+
+    Parameters
+    ----------
+    model : UNetModel
+        Model to validate.
+
+    Returns
+    -------
+    UNetModel
+        Validated model.
+
+    Raises
+    ------
+    ValueError
+        If the model has different number of input and output channels.
+    """
+    if model.num_classes != model.in_channels:
+        raise ValueError(
+            "The algorithm requires the same number of input and output channels. "
+            "Make sure that `in_channels` and `num_classes` are equal."
+        )
+
+    return model

careamics/lightning/callbacks/progress_bar_callback.py

@@ -5,7 +5,7 @@ from typing import Union
 
 from pytorch_lightning import LightningModule, Trainer
 from pytorch_lightning.callbacks import TQDMProgressBar
-from tqdm import tqdm
+from tqdm.auto import tqdm
 
 
 class ProgressBarCallback(TQDMProgressBar):

careamics/lightning/train_data_module.py

@@ -2,7 +2,6 @@
 
 from pathlib import Path
 from typing import Any, Callable, Literal, Optional, Union
-from warnings import warn
 
 import numpy as np
 import pytorch_lightning as L
@@ -261,11 +260,6 @@ class TrainDataModule(L.LightningDataModule):
 
         self.extension_filter: str = extension_filter
 
-        # Pytorch dataloader parameters
-        self.dataloader_params: dict[str, Any] = (
-            data_config.dataloader_params if data_config.dataloader_params else {}
-        )
-
     def prepare_data(self) -> None:
         """
         Hook used to prepare the data before calling `setup`.
@@ -447,21 +441,17 @@ class TrainDataModule(L.LightningDataModule):
         Any
             Training dataloader.
         """
-        # check because iterable dataset cannot be shuffled
-        if not isinstance(self.train_dataset, IterableDataset):
-            if ("shuffle" in self.dataloader_params) and (
-                not self.dataloader_params["shuffle"]
-            ):
-                warn(
-                    "Dataloader parameters include `shuffle=False`, this will be "
-                    "passed to the training dataloader and may result in bad results.",
-                    stacklevel=1,
-                )
-            else:
-                self.dataloader_params["shuffle"] = True
+        train_dataloader_params = self.data_config.train_dataloader_params.copy()
+
+        # NOTE: When next-gen datasets are completed this can be removed
+        # iterable dataset cannot be shuffled
+        if isinstance(self.train_dataset, IterableDataset):
+            del train_dataloader_params["shuffle"]
 
         return DataLoader(
-            self.train_dataset, batch_size=self.batch_size, **self.dataloader_params
+            self.train_dataset,
+            batch_size=self.batch_size,
+            **train_dataloader_params,
         )
 
     def val_dataloader(self) -> Any:
@@ -476,6 +466,7 @@ class TrainDataModule(L.LightningDataModule):
         return DataLoader(
             self.val_dataset,
             batch_size=self.batch_size,
+            **self.data_config.val_dataloader_params,
         )