careamics 0.0.15__py3-none-any.whl → 0.0.17__py3-none-any.whl

careamics/careamist.py

@@ -41,6 +41,7 @@ logger = get_logger(__name__)
 LOGGER_TYPES = list[Union[TensorBoardLogger, WandbLogger, CSVLogger]]
 
 
+# TODO type ignore have been added because of the czi data type in data configuration
 class CAREamist:
     """Main CAREamics class, allowing training and prediction using various algorithms.
 
@@ -208,17 +209,11 @@ class CAREamist:
 
         # instantiate trainer
         self.trainer = Trainer(
-            max_epochs=self.cfg.training_config.num_epochs,
-            precision=self.cfg.training_config.precision,
-            max_steps=self.cfg.training_config.max_steps,
-            check_val_every_n_epoch=self.cfg.training_config.check_val_every_n_epoch,
             enable_progress_bar=enable_progress_bar,
-            accumulate_grad_batches=self.cfg.training_config.accumulate_grad_batches,
-            gradient_clip_val=self.cfg.training_config.gradient_clip_val,
-            gradient_clip_algorithm=self.cfg.training_config.gradient_clip_algorithm,
             callbacks=self.callbacks,
             default_root_dir=self.work_dir,
             logger=experiment_logger,
+            **self.cfg.training_config.lightning_trainer_config or {},
         )
 
         # place holder for the datamodules
@@ -264,7 +259,7 @@ class CAREamist:
                 HyperParametersCallback(self.cfg),
                 ModelCheckpoint(
                     dirpath=self.work_dir / Path("checkpoints"),
-                    filename=self.cfg.experiment_name,
+                    filename=f"{self.cfg.experiment_name}_{{epoch:02d}}_step_{{step}}",
                     **self.cfg.training_config.checkpoint_callback.model_dump(),
                 ),
             ]
@@ -680,7 +675,7 @@ class CAREamist:
         # create the prediction
         self.pred_datamodule = create_predict_datamodule(
             pred_data=source,
-            data_type=data_type or self.cfg.data_config.data_type,
+            data_type=data_type or self.cfg.data_config.data_type,  # type: ignore
             axes=axes or self.cfg.data_config.axes,
             image_means=self.cfg.data_config.image_means,
             image_stds=self.cfg.data_config.image_stds,
@@ -826,11 +821,13 @@ class CAREamist:
         source_data_type: Literal["array", "tiff", "custom"]
         if isinstance(source, PredictDataModule):
             source_path = source.pred_data
-            source_data_type = source.data_type
+            source_data_type = source.data_type  # type: ignore
             extension_filter = source.extension_filter
-        elif isinstance(source, str | Path):
+        elif isinstance(source, (str | Path)):
             source_path = source
-            source_data_type = data_type or self.cfg.data_config.data_type
+            source_data_type = (
+                data_type or self.cfg.data_config.data_type  # type: ignore
+            )
             extension_filter = SupportedData.get_extension_pattern(
                 SupportedData(source_data_type)
             )
@@ -841,7 +838,7 @@ class CAREamist:
             raise ValueError(
                 "Predicting to disk is not supported for input type 'array'."
             )
-        assert isinstance(source_path, str | Path)  # because data_type != "array"
+        assert isinstance(source_path, (Path | str))  # because data_type != "array"
         source_path = Path(source_path)
 
         file_paths = list_files(source_path, source_data_type, extension_filter)
@@ -879,7 +876,7 @@ class CAREamist:
 
     def export_to_bmz(
         self,
-        path_to_archive: Union[Path, str],
+        path_to_archive: Union[Path | str],
         friendly_model_name: str,
         input_array: NDArray,
         authors: list[dict],

careamics/cli/conf.py

@@ -116,7 +116,11 @@ def care(  # numpydoc ignore=PR01
         ),
     ],
     batch_size: Annotated[int, typer.Option(help="Batch size.")],
-    num_epochs: Annotated[int, typer.Option(help="Number of epochs.")],
+    num_epochs: Annotated[int, typer.Option(help="Number of epochs.")] = 100,
+    num_steps: Annotated[
+        int | None,
+        typer.Option(help="Number of batches per epoch (limit_train_batches)."),
+    ] = None,
     data_type: Annotated[
         click.Choice,
         typer.Option(click_type=click.Choice(["tiff"]), help="Type of the data."),
@@ -175,6 +179,7 @@ def care(  # numpydoc ignore=PR01
         patch_size=patch_size,
         batch_size=batch_size,
         num_epochs=num_epochs,
+        num_steps=num_steps,
         # TODO: fix choosing augmentations
         augmentations=None if use_augmentations else [],
         independent_channels=independent_channels,
@@ -203,7 +208,11 @@ def n2n(  # numpydoc ignore=PR01
         ),
     ],
     batch_size: Annotated[int, typer.Option(help="Batch size.")],
-    num_epochs: Annotated[int, typer.Option(help="Number of epochs.")],
+    num_epochs: Annotated[int, typer.Option(help="Number of epochs.")] = 100,
+    num_steps: Annotated[
+        int | None,
+        typer.Option(help="Number of batches per epoch (limit_train_batches)."),
+    ] = None,
     data_type: Annotated[
         click.Choice,
         typer.Option(click_type=click.Choice(["tiff"]), help="Type of the data."),
@@ -259,6 +268,7 @@ def n2n(  # numpydoc ignore=PR01
         patch_size=patch_size,
         batch_size=batch_size,
         num_epochs=num_epochs,
+        num_steps=num_steps,
         # TODO: fix choosing augmentations
         augmentations=None if use_augmentations else [],
         independent_channels=independent_channels,
@@ -287,7 +297,11 @@ def n2v(  # numpydoc ignore=PR01
         ),
     ],
     batch_size: Annotated[int, typer.Option(help="Batch size.")],
-    num_epochs: Annotated[int, typer.Option(help="Number of epochs.")],
+    num_epochs: Annotated[int, typer.Option(help="Number of epochs.")] = 100,
+    num_steps: Annotated[
+        int | None,
+        typer.Option(help="Number of batches per epoch (limit_train_batches)."),
+    ] = None,
     data_type: Annotated[
         click.Choice,
         typer.Option(click_type=click.Choice(["tiff"]), help="Type of the data."),
@@ -364,6 +378,7 @@ def n2v(  # numpydoc ignore=PR01
         patch_size=patch_size,
         batch_size=batch_size,
         num_epochs=num_epochs,
+        num_steps=num_steps,
         # TODO: fix choosing augmentations
         augmentations=None if use_augmentations else [],
         independent_channels=independent_channels,

careamics/config/__init__.py

@@ -12,8 +12,10 @@ __all__ = [
     "Configuration",
     "DataConfig",
     "GaussianMixtureNMConfig",
+    "HDNAlgorithm",
     "InferenceConfig",
     "LVAELossConfig",
+    "MicroSplitAlgorithm",
     "MultiChannelNMConfig",
     "N2NAlgorithm",
     "N2VAlgorithm",
@@ -22,6 +24,8 @@ __all__ = [
     "VAEBasedAlgorithm",
     "algorithm_factory",
     "create_care_configuration",
+    "create_hdn_configuration",
+    "create_microsplit_configuration",
     "create_n2n_configuration",
     "create_n2v_configuration",
     "load_configuration",
@@ -30,6 +34,8 @@ __all__ = [
 
 from .algorithms import (
     CAREAlgorithm,
+    HDNAlgorithm,
+    MicroSplitAlgorithm,
     N2NAlgorithm,
     N2VAlgorithm,
     UNetBasedAlgorithm,
@@ -40,6 +46,8 @@ from .configuration import Configuration
 from .configuration_factories import (
     algorithm_factory,
     create_care_configuration,
+    create_hdn_configuration,
+    create_microsplit_configuration,
     create_n2n_configuration,
     create_n2v_configuration,
 )

careamics/config/algorithms/__init__.py

@@ -2,6 +2,8 @@
 
 __all__ = [
     "CAREAlgorithm",
+    "HDNAlgorithm",
+    "MicroSplitAlgorithm",
     "N2NAlgorithm",
     "N2VAlgorithm",
     "UNetBasedAlgorithm",
@@ -9,6 +11,8 @@ __all__ = [
 ]
 
 from .care_algorithm_model import CAREAlgorithm
+from .hdn_algorithm_model import HDNAlgorithm
+from .microsplit_algorithm_model import MicroSplitAlgorithm
 from .n2n_algorithm_model import N2NAlgorithm
 from .n2v_algorithm_model import N2VAlgorithm
 from .unet_algorithm_model import UNetBasedAlgorithm

careamics/config/algorithms/hdn_algorithm_model.py

@@ -0,0 +1,103 @@
+"""HDN algorithm configuration."""
+
+from typing import Literal
+
+from bioimageio.spec.generic.v0_3 import CiteEntry
+from pydantic import ConfigDict
+
+from careamics.config.algorithms.vae_algorithm_model import VAEBasedAlgorithm
+from careamics.config.architectures import LVAEModel
+from careamics.config.loss_model import LVAELossConfig
+
+HDN = "Hierarchical DivNoising"
+
+HDN_DESCRIPTION = (
+    "HDN leverages a hierarchical VAE to perform image "
+    "restoration. It is designed to be interpretable and unsupervised, "
+    "making it suitable for a wide range of microscopy images."
+)
+HDN_REF = CiteEntry(
+    text='Prakash, M., Delbracio, M., Milanfar, P., Jug, F. 2022. "Interpretable '
+    'Unsupervised Diversity Denoising and Artefact Removal." The International '
+    "Conference on Learning Representations (ICLR).",
+    doi="10.1561/2200000056",
+)
+
+
+class HDNAlgorithm(VAEBasedAlgorithm):
+    """HDN algorithm configuration."""
+
+    model_config = ConfigDict(validate_assignment=True)
+
+    algorithm: Literal["hdn"] = "hdn"
+
+    loss: LVAELossConfig
+
+    model: LVAEModel  # TODO add validators
+
+    is_supervised: bool = False
+
+    def get_algorithm_friendly_name(self) -> str:
+        """
+        Get the algorithm friendly name.
+
+        Returns
+        -------
+        str
+            Friendly name of the algorithm.
+        """
+        return HDN
+
+    def get_algorithm_keywords(self) -> list[str]:
+        """
+        Get algorithm keywords.
+
+        Returns
+        -------
+        list[str]
+            List of keywords.
+        """
+        return [
+            "restoration",
+            "VAE",
+            "3D" if self.model.is_3D() else "2D",
+            "CAREamics",
+            "pytorch",
+        ]
+
+    def get_algorithm_references(self) -> str:
+        """
+        Get the algorithm references.
+
+        This is used to generate the README of the BioImage Model Zoo export.
+
+        Returns
+        -------
+        str
+            Algorithm references.
+        """
+        return HDN_REF.text + " doi: " + HDN_REF.doi
+
+    def get_algorithm_citations(self) -> list[CiteEntry]:
+        """
+        Return a list of citation entries of the current algorithm.
+
+        This is used to generate the model description for the BioImage Model Zoo.
+
+        Returns
+        -------
+        List[CiteEntry]
+            List of citation entries.
+        """
+        return [HDN_REF]
+
+    def get_algorithm_description(self) -> str:
+        """
+        Get the algorithm description.
+
+        Returns
+        -------
+        str
+            Algorithm description.
+        """
+        return HDN_DESCRIPTION

careamics/config/algorithms/microsplit_algorithm_model.py

@@ -0,0 +1,103 @@
+"""MicroSplit algorithm configuration."""
+
+from typing import Literal
+
+from bioimageio.spec.generic.v0_3 import CiteEntry
+from pydantic import ConfigDict
+
+from careamics.config.algorithms.vae_algorithm_model import VAEBasedAlgorithm
+from careamics.config.architectures import LVAEModel
+from careamics.config.loss_model import LVAELossConfig
+
+MICROSPLIT = "MicroSplit"
+
+MICROSPLIT_DESCRIPTION = """MicroSplit is a self-supervised deep learning method for
+microscopy image splitting that combines the strengths of both denoising and
+representation learning approaches."""
+
+MICROSPLIT_REF = CiteEntry(
+    text='Prakash, M., Delbracio, M., Milanfar, P., Jug, F. 2022. "Interpretable '
+    'Unsupervised Diversity Denoising and Artefact Removal." The International '
+    "Conference on Learning Representations (ICLR).",
+    doi="10.1561/2200000056",
+)
+
+
+class MicroSplitAlgorithm(VAEBasedAlgorithm):
+    """MicroSplit algorithm configuration."""
+
+    model_config = ConfigDict(validate_assignment=True)
+
+    algorithm: Literal["microsplit"] = "microsplit"
+
+    loss: LVAELossConfig
+
+    model: LVAEModel  # TODO add validators
+
+    is_supervised: bool = True
+
+    def get_algorithm_friendly_name(self) -> str:
+        """
+        Get the algorithm friendly name.
+
+        Returns
+        -------
+        str
+            Friendly name of the algorithm.
+        """
+        return MICROSPLIT
+
+    def get_algorithm_keywords(self) -> list[str]:
+        """
+        Get algorithm keywords.
+
+        Returns
+        -------
+        list[str]
+            List of keywords.
+        """
+        return [
+            "restoration",
+            "VAE",
+            "self-supervised",
+            "3D" if self.model.is_3D() else "2D",
+            "CAREamics",
+            "pytorch",
+        ]
+
+    def get_algorithm_references(self) -> str:
+        """
+        Get the algorithm references.
+
+        This is used to generate the README of the BioImage Model Zoo export.
+
+        Returns
+        -------
+        str
+            Algorithm references.
+        """
+        return MICROSPLIT_REF.text + " doi: " + MICROSPLIT_REF.doi
+
+    def get_algorithm_citations(self) -> list[CiteEntry]:
+        """
+        Return a list of citation entries of the current algorithm.
+
+        This is used to generate the model description for the BioImage Model Zoo.
+
+        Returns
+        -------
+        List[CiteEntry]
+            List of citation entries.
+        """
+        return [MICROSPLIT_REF]
+
+    def get_algorithm_description(self) -> str:
+        """
+        Get the algorithm description.
+
+        Returns
+        -------
+        str
+            Algorithm description.
+        """
+        return MICROSPLIT_DESCRIPTION

careamics/config/algorithms/n2v_algorithm_model.py

@@ -1,10 +1,9 @@
 """N2V Algorithm configuration."""
 
-from typing import Annotated, Literal
+from typing import Annotated, Literal, Self
 
 from bioimageio.spec.generic.v0_3 import CiteEntry
 from pydantic import AfterValidator, ConfigDict, model_validator
-from typing_extensions import Self
 
 from careamics.config.architectures import UNetModel
 from careamics.config.support import SupportedPixelManipulation, SupportedStructAxis

careamics/config/algorithms/vae_algorithm_model.py

@@ -3,10 +3,9 @@
 from __future__ import annotations
 
 from pprint import pformat
-from typing import Literal
+from typing import Literal, Self
 
 from pydantic import BaseModel, ConfigDict, model_validator
-from typing_extensions import Self
 
 from careamics.config.architectures import LVAEModel
 from careamics.config.likelihood_model import (
@@ -40,14 +39,17 @@ class VAEBasedAlgorithm(BaseModel):
     # defined in SupportedAlgorithm
     # TODO: Use supported Enum classes for typing?
     #   - values can still be passed as strings and they will be cast to Enum
-    algorithm: Literal["musplit", "denoisplit"]
+    algorithm: Literal["hdn", "microsplit"]
 
     # NOTE: these are all configs (pydantic models)
     loss: LVAELossConfig
     model: LVAEModel
     noise_model: MultiChannelNMConfig | None = None
     noise_model_likelihood: NMLikelihoodConfig | None = None
-    gaussian_likelihood: GaussianLikelihoodConfig | None = None
+    gaussian_likelihood: GaussianLikelihoodConfig | None = None  # TODO change to str
+
+    mmse_count: int = 1
+    is_supervised: bool = False
 
     # Optional fields
     optimizer: OptimizerModel = OptimizerModel()
@@ -64,22 +66,26 @@ class VAEBasedAlgorithm(BaseModel):
         Self
             The validated model.
         """
-        # musplit
-        if self.algorithm == SupportedAlgorithm.MUSPLIT:
-            if self.loss.loss_type != SupportedLoss.MUSPLIT:
+        # hdn
+        # TODO move to designated configurations
+        if self.algorithm == SupportedAlgorithm.HDN:
+            if self.loss.loss_type != SupportedLoss.HDN:
                 raise ValueError(
-                    f"Algorithm {self.algorithm} only supports loss `musplit`."
+                    f"Algorithm {self.algorithm} only supports loss `hdn`."
                 )
-
-        if self.algorithm == SupportedAlgorithm.DENOISPLIT:
+            if self.model.multiscale_count > 1:
+                raise ValueError("Algorithm `hdn` does not support multiscale models.")
+        # musplit
+        if self.algorithm == SupportedAlgorithm.MICROSPLIT:
             if self.loss.loss_type not in [
+                SupportedLoss.MUSPLIT,
                 SupportedLoss.DENOISPLIT,
                 SupportedLoss.DENOISPLIT_MUSPLIT,
-            ]:
+            ]:  # TODO Update losses configs, make loss just microsplit
                 raise ValueError(
-                    f"Algorithm {self.algorithm} only supports loss `denoisplit` "
-                    "or `denoisplit_musplit."
-                )
+                    f"Algorithm {self.algorithm} only supports loss `microsplit`."
+                )  # TODO Update losses configs
+
             if (
                 self.loss.loss_type == SupportedLoss.DENOISPLIT
                 and self.model.predict_logvar is not None
@@ -88,8 +94,10 @@ class VAEBasedAlgorithm(BaseModel):
                     "Algorithm `denoisplit` with loss `denoisplit` only supports "
                     "`predict_logvar` as `None`."
                 )
-
-            if self.noise_model is None:
+            if (
+                self.loss.loss_type == SupportedLoss.DENOISPLIT
+                and self.noise_model is None
+            ):
                 raise ValueError("Algorithm `denoisplit` requires a noise model.")
         # TODO: what if algorithm is not musplit or denoisplit
         return self
@@ -108,6 +116,12 @@ class VAEBasedAlgorithm(BaseModel):
                 f"Number of output channels ({self.model.output_channels}) must match "
                 f"the number of noise models ({len(self.noise_model.noise_models)})."
             )
+
+        if self.algorithm == SupportedAlgorithm.HDN:
+            assert self.model.output_channels == 1, (
+                f"Number of output channels ({self.model.output_channels}) must be 1 "
+                "for algorithm `hdn`."
+            )
         return self
 
     @model_validator(mode="after")
@@ -127,6 +141,16 @@ class VAEBasedAlgorithm(BaseModel):
                 "Gaussian likelihood model `predict_logvar` "
                 f"({self.gaussian_likelihood.predict_logvar}).",
             )
+        # if self.algorithm == SupportedAlgorithm.HDN:
+        #     assert (
+        #         self.model.predict_logvar is None
+        #     ), "Model `predict_logvar` must be `None` for algorithm `hdn`."
+        #     if self.gaussian_likelihood is not None:
+        #         assert self.gaussian_likelihood.predict_logvar is None, (
+        #             "Gaussian likelihood model `predict_logvar` must be `None` "
+        #             "for algorithm `hdn`."
+        #         )
+        # TODO check this
         return self
 
     def __str__(self) -> str:
@@ -138,3 +162,14 @@ class VAEBasedAlgorithm(BaseModel):
             Pretty string.
         """
         return pformat(self.model_dump())
+
+    @classmethod
+    def get_compatible_algorithms(cls) -> list[str]:
+        """Get the list of compatible algorithms.
+
+        Returns
+        -------
+        list of str
+            List of compatible algorithms.
+        """
+        return ["hdn", "microsplit"]

careamics/config/architectures/lvae_model.py

@@ -1,9 +1,8 @@
 """LVAE Pydantic model."""
 
-from typing import Literal
+from typing import Literal, Self
 
 from pydantic import ConfigDict, Field, field_validator, model_validator
-from typing_extensions import Self
 
 from .architecture_model import ArchitectureModel
 
@@ -15,11 +14,9 @@ class LVAEModel(ArchitectureModel):
     model_config = ConfigDict(validate_assignment=True, validate_default=True)
 
     architecture: Literal["LVAE"]
-    """Name of the architecture."""
-
-    input_shape: list[int] = Field(default=[64, 64], validate_default=True)
-    """Shape of the input patch (C, Z, Y, X) or (C, Y, X) if the data is 2D."""
 
+    input_shape: tuple[int, ...] = Field(default=(64, 64), validate_default=True)
+    """Shape of the input patch (Z, Y, X) or (Y, X) if the data is 2D."""
     encoder_conv_strides: list = Field(default=[2, 2], validate_default=True)
 
     # TODO make this per hierarchy step ?
@@ -42,7 +39,7 @@ class LVAEModel(ArchitectureModel):
         default="ELU",
     )
 
-    predict_logvar: Literal[None, "pixelwise"] = None
+    predict_logvar: Literal[None, "pixelwise"] = "pixelwise"
     analytical_kl: bool = Field(default=False)
 
     @model_validator(mode="after")
@@ -126,6 +123,13 @@ class LVAEModel(ArchitectureModel):
                 f"Input shape must be greater than 1 in all dimensions"
                 f"(got {input_shape})."
             )
+
+        if any(s < 64 for s in input_shape[-2:]):
+            raise ValueError(
+                f"Input shape must be greater or equal to 64 in XY dimensions"
+                f"(got {input_shape})."
+            )
+
         return input_shape
 
     @field_validator("encoder_n_filters")
@@ -255,4 +259,4 @@ class LVAEModel(ArchitectureModel):
         bool
             Whether the model is 3D or not.
         """
-        return self.conv_dims == 3
+        return len(self.input_shape) == 3

careamics/config/callback_model.py

@@ -24,7 +24,7 @@ class CheckpointModel(BaseModel):
 
     model_config = ConfigDict(validate_assignment=True, validate_default=True)
 
-    monitor: Literal["val_loss"] = Field(default="val_loss")
+    monitor: Literal["val_loss"] | str | None = Field(default="val_loss")
     """Quantity to monitor, currently only `val_loss`."""
 
     verbose: bool = Field(default=False)
@@ -36,8 +36,12 @@ class CheckpointModel(BaseModel):
     save_last: Literal[True, False, "link"] | None = Field(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=100)
-    """If `save_top_k == kz, the best k models according to the quantity monitored
+    save_top_k: int = Field(
+        default=3,
+        ge=-1,
+        le=100,
+    )
+    """If `save_top_k == k, 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."""